Python Function 重复使用

Question

I have a function defined as below in my python code:

def insert_source():
    sql = "INSERT INTO public.source (width, height, sha1) " \
          "VALUES (%s, %s, %s)"
    conn = None
    data = widtho, heighto, sha1o
    try:
        # read database configuration
        params = config()
        # connect to the PostgreSQL database
        conn = psycopg2.connect(**params)
        # create a new cursor
        cur = conn.cursor()
        # execute the INSERT statement
        cur.execute(sql, data)
        # commit the changes to the database
        conn.commit()
        print('Source image data successfully imported', sha1o)
        # close communication with the database
        cur.close()
    except (Exception, psycopg2.DatabaseError) as error:
        if str(error) in "duplicate key value violates unique constraint":
            raise
        else:
            print("I've seen this image before, skipping insert")
    finally:
        if conn is not None:
            conn.close()
            return id

if __name__ == '__main__':
    insert_source()

this works fine and inserts that or throws an exception as expected. BUT now later in my script I want to write some more data in a different table within the same db. Should I copy/paste the entire function above and rename it or is there a better way to say do most of these things but the sql and data should be X?

FWIW duplicating that function results in a Duplicated code fragment (31 lines long) note so it just seems silly to replicate 98% of that code when I need to change 4 lines in total.

Answer 1

you can add arguments to the function for example table_name, things_to_insert, and values.. so your function will be like:

def insert_source(table_name, things_to_insert, values):
    sql = f"INSERT INTO {table_name} {things_to_insert} VALUES {values}"
    conn = None
    data = widtho, heighto, sha1o
    try:
        # read database configuration
        params = config()
        # connect to the PostgreSQL database
        conn = psycopg2.connect(**params)
        # create a new cursor
        cur = conn.cursor()
        # execute the INSERT statement
        cur.execute(sql, data)
        # commit the changes to the database
        conn.commit()
        print('Source image data successfully imported', sha1o)
        # close communication with the database
        cur.close()
    except (Exception, psycopg2.DatabaseError) as error:
        if str(error) in "duplicate key value violates unique constraint":
            raise
        else:
            print("I've seen this image before, skipping insert")
    finally:
        if conn is not None:
            conn.close()
            return id

if __name__ == '__main__':
    table_name = 'public.source'
    thing_to_insert = (width, height, sha1)
    values = '(%s, %s, %s)'
    insert_source(table_name, things_to_insert, values)

this will do exactly what your code does, and you can customize arguments to make the function reusable.

BUT, as Ted Klein Bergman said in the comments, this way is great candidate for SQL injections. so what you have to do is to make prepared/parametrized queries as Klein suggested.