How do I pass arguments to my pl/python custom function

Question

I am trying to create my first ever custom function that uses arguments in pl/python. I have created triggers before using pl/python language but these kind of functions don't have arguments. The code of the function I have created is below:

CREATE OR REPLACE FUNCTION test_pl_pyton_func(table_name TEXT, field_name TEXT) RETURNS TEXT AS
    $BODY$
    plan = plpy.prepare("SELECT DISTINCT $1 FROM $2", ['TEXT', 'TEXT'])
    result = plpy.execute(plan, [field_name, table_name])
    while i < result.nrows():
        print(result[i][field_name])
        i = i+1
    $BODY$
    LANGUAGE plpython3u;

Now if I run:

SELECT test_pl_pyton_func(table_name, field_name);

I get the following error:

spiexceptions.SyntaxError: syntax error at or near "$2"
    LINE 1: SELECT DISTINCT $1 FROM $2
                                    ^
        QUERY:  SELECT DISTINCT $1 FROM $2
        CONTEXT:  Traceback (most recent call last):
          PL/Python function "test_pl_pyton_func", line 2, in <module>
            plan = plpy.prepare("SELECT DISTINCT $1 FROM $2", ['TEXT', 'TEXT'])
        PL/Python function "test_pl_pyton_func"

What am I getting wrong?

Answer 1

You cannot put neither table name nor field names as prepared statement's arguments. The table names and the queried columns must always be constant in a prepared query. If they weren't, the DBMS wouldn't know the return types for the query, so it wouldn't be able to prepare the query.

For your use case you can use simple python formatting combined with plpy's quoting (escaping):

plpy.prepare(f"SELECT DISTINCT {plpy.quote_ident(field_name)} FROM {plpy.quote_ident(table_name)}")

More on the escaping functions here: https://www.postgresql.org/docs/current/plpython-util.html (see: plpy.quote_literal, plpy.quote_ident, plpy.quote_nullable).