How would I do a bulk update using Python 3 | psycopg 3.1.8?

Question

In python I have a list of tuples

Each tuple has four fields, the first field is the row identifier that would be used in a where clause in a SQL update statement.

tuple: (id, status, department, manager_email)

update_tuple_list = [
 (1, 'ACTIVE', 'UC563', '[email protected]'),
 (2, 'ACTIVE', 'UC921', '[email protected]'),
 (3, 'DISABLED', 'UC983', '[email protected]'),
  ...
]

Using Python 3, psycopg ver 3.1.8

How would I do a bulk update?

For table user, the statement would be:

update user set
    status = %s,
    department = %s,
    manager_email = %s
    where id  %s;   # where the value of 'id' is the first field in the tuple.

I can rearrange the tuple to put the row id id at the end of the tuple if that would help.

How do I feed the example list of tuples to a statement like the above and do a single bulk update?

---

Is there a better more understandable way to do this using Python 3 | psycopg Ver 3.1.8 (NOT psycopg2)

I have seen stack overflow answers involving psycopg2 and I simply do not understand the response or the solution provided there, and also I am not using psycopg2. It's just not clear.

It's very simple with a bulk insert because there is no where clause'. With a bulk update, there is a where clause to identify the row that is being updated.

I have searched on this and tried to find a simple solution using the latest python version and library version, and nothing comes up.

Any and all help to solve this would be greatly appreciated.

Answer 1

\d animals
                        Table "public.animals"
   Column   |          Type          | Collation | Nullable | Default 
------------+------------------------+-----------+----------+---------
 pk_animals | integer                |           | not null | 
 cond       | character varying(200) |           | not null | 
 animal     | character varying(200) |           | not null | 

select * from animals where pk_animals in (3, 16);

 pk_animals | cond  | animal  
------------+-------+---------
         16 | fair  | heron
          3 | good  | mole

import psycopg
con = psycopg.connect("dbname=test host=localhost  user=postgres")
cur = con.cursor()

a_list = [('horse', 'fair', 16), ('lion', 'good', 3)]

cur.executemany('update animals set (animal, cond) = (%s, %s) where pk_animals = %s', a_list)

con.commit()

select * from animals where pk_animals in (3, 16);
 pk_animals | cond | animal 
------------+------+--------
         16 | fair | horse
          3 | good | lion

Answer 2

How about this ?:

import pandas as pd
import psycopg
import logging
import io

logger = logging.getLogger(__name__)

map_df2pg = {
    'int64': 'int8',
    'int16': 'int2',
    'int32': 'int4',
    'object': 'text',
    'string': 'text',
    'boolean': 'bool',
    'datetime64[ns]': 'timestamp',
    'float32': 'float4',
    'float64': 'float8'
}


def execute(query: str, params=None, conn: psycopg.Connection = None) -> psycopg.Cursor:
    try:
        if conn is None:
            with psycopg.connect(your_db_uri) as conn:
                cur = conn.execute(query=query, params=params)
        else:
            cur = conn.execute(query=query, params=params)
    except Exception as e:
        logger.exception('Failed query execution', e)
        raise e
    return cur


def get_column_names(table_name: str, conn: psycopg.Connection = None):
    if conn is None:
        with psycopg.connect(your_db_uri) as conn:
            cursor = conn.execute(f"SELECT * FROM {table_name} LIMIT 0")
            col_names = [desc[0] for desc in cursor.description]
    else:
        cursor = conn.execute(f"SELECT * FROM {table_name} LIMIT 0")
        col_names = [desc[0] for desc in cursor.description]
    return col_names


def insert(df, table_name, buffer_size: int = 1024, conn: psycopg.Connection = None):
    cols = df.columns.tolist()

    def _insert_bulk(conn: psycopg.Connection):
        cursor = conn.cursor()
        cols_in_db = get_column_names(table_name, conn=conn)
        _df = df.reindex(columns=cols_in_db)  # must match the natural order of the db
        buffer = io.StringIO()
        _df.to_csv(buffer, index=False)
        buffer.seek(0)
        with cursor.copy(f'copy {table_name} from stdin with (format csv, header)') as copy:
            while data := buffer.read(buffer_size):
                copy.write(data)

    if conn is None:
        with psycopg.connect(your_db_uri) as conn:
            _insert_bulk(conn)
    else:
        _insert_bulk(conn)


def update(df: pd.DataFrame, table_name: str, on_cols: list = []):
    cols = df.columns.tolist()
    col_types = [d.name for d in df.dtypes]
    update_cols = list(set(cols) - set(on_cols))

    _cols = update_cols + on_cols  # place on cols at the end
    _df = df.reindex(columns=_cols)

    has_where_clause = False
    if len(on_cols) > 0:
        has_where_clause = True

    with psycopg.connect(your_db_uri) as conn:
        # create an empty temporary table
        cols_as_type = ", ".join([f'{col} {map_df2pg[col_type]}' for col, col_type in zip(cols, col_types)])
        query_tmp_tbl = f"create temp table tmptbl({cols_as_type}) on commit drop"
        execute(query_tmp_tbl, conn=conn)

        # insert df data into temporary table
        insert(df, table_name='tmptbl', conn=conn)

        # update table from temporary table
        update_query = f"""
            UPDATE {table_name} b
            SET {', '.join([f"{col} = a.{col}" for col in update_cols])}
            FROM tmptbl a
            """

        if has_where_clause:
            where_clause = ' AND '.join([f"a.{col} = b.{col}" for col in on_cols])
            update_query = f"{update_query} WHERE {where_clause}"

        execute(update_query, conn=conn)

Instead of looping, you first create a temporary table with your new data and update your target table from it in one go. This has served me well and is super fast.