How can a PostgreSQL function accept an integer or array of integers for the same function?

Question

I have a function in Postgres 9.6 that accepts an int[] parameter. I'd like to have the function accept a single int as well (and convert it to a single element array if necessary).

 CREATE OR REPLACE FUNCTION get_subordinates(inp_persona_ids integer[])
 -- Get all subordnates of the people passed in as array
 -- TODO allow a single persona ID (int) to be passed in as inp_persona_ids
 RETURNS TABLE (persona_id int) AS
 $$
    BEGIN
        RETURN QUERY(
            WITH RECURSIVE children AS (
                -- passed in persona_id
                SELECT
                    id AS persona_id,
                    manager_id
                FROM
                    personas
                WHERE
                    id = ANY(inp_persona_ids)
                UNION
                -- and all subordinates
                SELECT
                    p.id AS persona_id,
                    p.manager_id
                FROM
                    personas p
                    JOIN children c ON p.manager_id = c.persona_id
            )
            SELECT 
                children.persona_id
            FROM
                children
            LEFT JOIN
                personas on children.persona_id = personas.id
            WHERE personas.disabled IS NOT TRUE
        );
    END;
$$ LANGUAGE plpgsql

How would I change the function definition and also add some conditional logic to test for int and change to ARRAY[int] if necessary?

score 2 · Accepted Answer · answered Sep 23 '19 at 01:24

It is not possible to handle this in a single function, but you can just overload the function with an integer parameter and pass this as an array to your existing function:

CREATE OR REPLACE FUNCTION get_subordinates(inp_persona_id integer)
RETURNS TABLE (persona_id int) AS
$$
BEGIN
    RETURN QUERY SELECT * FROM get_subordinates(ARRAY[inp_persona_id]);
END;
$$ LANGUAGE plpgsql;

Perhaps you might also want to check the argument(s) against NULL, this is up to you.

score 1 · Answer 2 · answered Sep 25 '19 at 22:58

It is possible with a single function using the VARIADIC modifier:

CREATE OR REPLACE FUNCTION get_subordinates(VARIADIC inp_persona_ids int[])
  RETURNS TABLE (persona_id int) AS
$func$
WITH RECURSIVE children AS (   -- passed in persona_id
   SELECT id AS persona_id, manager_id, disabled
   FROM   personas
   WHERE  id = ANY(inp_persona_ids)

   UNION ALL   -- and all subordinates
   SELECT p.id AS persona_id
        , p.manager_id
   FROM   children c
   JOIN   personas p ON p.manager_id = c.persona_id
)
SELECT c.persona_id
FROM   children c
WHERE  c.disabled IS NOT TRUE
$func$  LANGUAGE sql;

But you need to add the keyword VARIADIC in the call when providing an array instead of a list:

SELECT * FROM  get_subordinates(VARIADIC '{1,2,3}'::int[]);
SELECT * FROM  get_subordinates(1,2,3);
SELECT * FROM  get_subordinates(1);

If that's not an option you are back to function overloading as suggested in another answer.

See:

Asides

Looks like this can be a simpler SQL function.
UNION made no sense. Duplicates can only occur if your tree goes in circles, which would create an endless loop and the rCTE would error out. Use the cheaper UNION ALL.
LEFT JOIN made no sense. The added WHERE forced it to behave like a plain [INNER] JOIN anyway.
But remove the join completely and retrieve the column disabled inside the rCTE.

Expecting the upstream caller to use variadic for arrays won't work, so the overloading is the solution here. The simplification in your example helps with readability and performance. Appreciate that. — Matt Graham, Sep 26 '19 at 22:16

How can a PostgreSQL function accept an integer or array of integers for the same function?

2 Answers2

Asides