如何在Python/PostgreSQL中批量插入时检查约束并返回主键

Absolutely! You can pull off batch inserts with constraint checks and get back all corresponding primary keys in a single query—here's how to do it right with PostgreSQL and psycopg2, working with your serial primary keys and unique constraints (no more slow single-row inserts or post-insert lookups).

Core Issue with Your Current Setup #

Your single-row insert works, but it’s inefficient because every row requires a round-trip to the database. Worse, the fallback lookup for existing IDs adds even more overhead. We can fix this by combining bulk insert syntax, ON CONFLICT handling, and psycopg2’s built-in batch tools to get all required IDs in one go.

Solution 1: Use ON CONFLICT DO UPDATE (No-Op) to Return All IDs #

Since ON CONFLICT DO NOTHING won’t return IDs for rows that hit a constraint, we can use a harmless "no-op" update (setting a field to itself) to trigger RETURNING for every row—whether it was inserted or already existed. This is the simplest approach.

Here’s the code using psycopg2’s execute_values (optimized for bulk operations):

from psycopg2.extras import execute_values

# Your existing variables (adjust as needed)
schema_name = "your_schema"
table_name = "your_table"
pk_col = head[0]  # e.g., 'id' (serial primary key)
insert_cols = head[1:]  # Columns you're inserting data into
unique_col = unique_field  # Column with your unique constraint

# Build the bulk insert SQL
sql = f"""
INSERT INTO {schema_name}.{table_name} ({','.join(insert_cols)})
VALUES %s
ON CONFLICT ({unique_col}) DO UPDATE 
SET {pk_col} = {pk_col}  # No-op update: doesn't change data, but triggers RETURNING
RETURNING {pk_col}, {unique_col};
"""

# Prepare your bulk data: list of tuples matching insert_cols
# Example: rows = [(val1, val2), (val3, val4), ...]
execute_values(pdb.cur, sql, rows)

# Build a mapping of unique values to their primary keys (new or existing)
id_mapping = {}
for pk_val, unique_val in pdb.cur.fetchall():
    id_mapping[unique_val] = pk_val

Solution 2: CTE-Based Approach (No Fake Updates) #

If you prefer to avoid even a no-op update, use a CTE to first insert new rows, then fetch IDs for existing rows, and combine the results:

from psycopg2.extras import execute_values

# Same variables as above
sql = f"""
WITH inserted AS (
    INSERT INTO {schema_name}.{table_name} ({','.join(insert_cols)})
    VALUES %s
    ON CONFLICT ({unique_col}) DO NOTHING
    RETURNING {pk_col}, {unique_col}
),
existing AS (
    SELECT {pk_col}, {unique_col}
    FROM {schema_name}.{table_name}
    WHERE {unique_col} IN (SELECT unnest(%s))
)
SELECT * FROM inserted
UNION ALL
SELECT * FROM existing;
"""

# Extract unique values from your bulk data to query existing rows
unique_values = [row[insert_cols.index(unique_col)] for row in rows]

# Execute with both the bulk data and unique values list
execute_values(pdb.cur, sql, rows, vars_list=(unique_values,))

# Build your ID mapping the same way
id_mapping = {}
for pk_val, unique_val in pdb.cur.fetchall():
    id_mapping[unique_val] = pk_val

Key Advantages Over Your Current Setup #

• Massive Performance Boost: Reduces database round-trips from N to 1 (where N is your number of rows)
• No SQL Injection Risk: execute_values safely handles parameterization, unlike manual string formatting
• Single Source of Truth: Gets all primary keys (new and existing) in one query—no need for fallback SELECT calls
• Scalable: Adjust the batch size (e.g., 1000-10000 rows per batch) based on your database's capacity

Important Notes #

• Make sure your ON CONFLICT clause matches exactly with your unique constraint (use (col1, col2) for composite unique keys)
• execute_values automatically optimizes the batch insert plan for PostgreSQL
• Test with your specific data size to find the optimal batch size (too small = still slow; too large = memory issues)

内容的提问来源于stack exchange，提问作者AnarKi