Using psycopg2 to insert, return and update in bulk

Question

I'm using postgres to log orders. I have one function for inserting a row for each order and another function that updates the order once the order has gone through and been confirmed. As of now I'm using 'INSERT' to insert a order and RETURN the row_id. I then use UPDATE with the row_id that was returned in the INSERT to update the order. How can I do this in bulk? I see that psycopg2 has a executemany function, but according to the documentation, it can't return anything. Is there any way to do this?

def initial_log(self, orders):
    with self.conn.cursor() as cur:
        row_ids = []
        for order in orders:
            cur.execute('INSERT INTO orders (order_num, order_amount)
                         VALUES (%s, %s) RETURNING order_id;' 
                         (order.num, order.amount))
            row_id = cur.fetchone()[0]
            row_ids.append(row_id)
        self.conn.commit()
        return row_ids


def update_log(self, row_ids, updated_orders):
    with self.conn.cursor() as cur:
        for row_id, order in zip(row_ids, updated_orders):
            status_list = order.messages
            encoded_status = encode_status(status_list)
            cur.execute('UPDATE orders SET (final_order_amount, order_ack,
                    order_time, status) = 
                    (%s, %s, current_timestamp, %s) WHERE order_id = %s',
                    (order.final_amount,
                     order.ack_num, encoded_status, row_id))
        self.conn.commit()

Clodoaldo Neto · Accepted Answer · 2016-02-14 21:43:19Z

def initial_log(self, orders):

    insert = '''
        insert into orders (order_num, order_amount)
        values {}
        returning order_id;
    '''.format(','.join(['%s'] * len(orders)))

    t = [(order.num, order.amount) for order in orders]

    cur = self.conn.cursor()
    # print cur.mogrify(insert, t)
    cur.execute(insert, t)
    rs = cur.fetchall()
    self.conn.commit()

    row_ids = [row[0] for row in rs]
    return row_ids

def update_log(self, row_ids, updated_orders):

    update = '''
        update orders o
        set (final_order_amount, order_ack, order_time, status) =
            (s.final_order_amount, s.order_ack, current_timestamp, s.status)
        from (values
            {}
        ) s (final_order_amount, order_ack, status, order_id)
        where o.order_id = s.order_id
    '''.format(','.join(['%s'] * len(updated_orders)))

    t = [
        (order.final_amount, order.ack_num, encode_status(order.messages), row_id)
        for row_id, order in zip(row_ids, updated_orders)
    ]

    cur = self.conn.cursor()
    # print cur.mogrify(update, t)
    cur.execute(update, t)
    self.conn.commit()

Uncomment cursor.mogrify when checking what is being sent to the server.

Collectives™ on Stack Overflow

Using psycopg2 to insert, return and update in bulk

1 Answer 1

Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Related