如何用Python/psycopg2实现PostgreSQL高性能插入Offer关联数据

首先修正原有表结构的不合理问题：现有Offers表的category字段定义为SERIAL类型不符合外键设计要求，SERIAL是自增整数类型，作为关联外键仅需使用INT即可，修改后的字段定义为：

category INT REFERENCES Categories (id) ON UPDATE CASCADE ON DELETE CASCADE

OfferRegions表的offer和region字段同理，也建议改为INT类型，避免不必要的自增属性冲突。同时需要给Offers表添加UNIQUE(title, category)约束，用于判断重复Offer。

核心优化思路 #

• 利用PostgreSQL唯一约束 + INSERT ... ON CONFLICT DO NOTHING语法，替代原有NOT EXISTS判断，避免重复查询
• 新增极少的分类、地区数据在Python端做内存缓存，避免每次插入都查询数据库
• 合并多步SQL操作为单次数据库往返，减少网络IO开销
• 批量场景下使用批量插入接口，进一步提升吞吐量

具体实现步骤 #

1. 预加载分类、地区缓存 #

程序启动时全量拉取已存在的分类和地区ID映射，后续新增时先查缓存，不存在再写入数据库并更新缓存：

import psycopg2
from psycopg2.extras import execute_values

# 初始化全局缓存
cat_cache = {}
region_cache = {}

def init_cache(conn):
    global cat_cache, region_cache
    with conn.cursor() as cur:
        cur.execute("SELECT name, id FROM Categories")
        cat_cache = {row[0]: row[1] for row in cur.fetchall()}
        cur.execute("SELECT name, id FROM Regions")
        region_cache = {row[0]: row[1] for row in cur.fetchall()}

2. 分类/地区不存在时自动写入并更新缓存 #

def get_cat_id(conn, cat_name):
    if cat_name in cat_cache:
        return cat_cache[cat_name]
    with conn.cursor() as cur:
        cur.execute("""
            INSERT INTO Categories(name) VALUES (%s)
            ON CONFLICT (name) DO NOTHING
            RETURNING id
        """, (cat_name,))
        res = cur.fetchone()
        if res:
            cat_id = res[0]
        else:
            # 并发场景下可能已经被其他进程插入，补充一次查询
            cur.execute("SELECT id FROM Categories WHERE name = %s", (cat_name,))
            cat_id = cur.fetchone()[0]
        cat_cache[cat_name] = cat_id
        return cat_id

def get_region_id(conn, region_name):
    if region_name in region_cache:
        return region_cache[region_name]
    with conn.cursor() as cur:
        cur.execute("""
            INSERT INTO Regions(name) VALUES (%s)
            ON CONFLICT (name) DO NOTHING
            RETURNING id
        """, (region_name,))
        res = cur.fetchone()
        if res:
            region_id = res[0]
        else:
            cur.execute("SELECT id FROM Regions WHERE name = %s", (region_name,))
            region_id = cur.fetchone()[0]
        region_cache[region_name] = region_id
        return region_id

3. 单条Offer插入优化 #

用CTE合并Offer和OfferRegions的插入操作，单次数据库往返完成全流程：

def insert_offer(conn, offer):
    cat_id = get_cat_id(conn, offer.category)
    # 先获取所有关联地区ID，缓存命中时无数据库查询
    region_ids = [get_region_id(conn, r) for r in offer.regions]
    with conn.cursor() as cur:
        cur.execute("""
            WITH inserted_offer AS (
                INSERT INTO Offers(title, category) VALUES (%s, %s)
                ON CONFLICT (title, category) DO NOTHING
                RETURNING id
            )
            INSERT INTO OfferRegions(offer, region)
            SELECT id, unnest(%s::INT[]) FROM inserted_offer
            ON CONFLICT (offer, region) DO NOTHING
        """, (offer.title, cat_id, region_ids))
    conn.commit()

4. 批量插入优化（多Offer场景） #

如果需要一次性插入多个Offer，使用execute_values实现批量写入，性能提升10倍以上：

def batch_insert_offers(conn, offers):
    # 先预处理所有分类和地区ID，全部走缓存基本无数据库查询
    offer_rows = []
    offer_region_rows = []
    for offer in offers:
        cat_id = get_cat_id(conn, offer.category)
        offer_rows.append((offer.title, cat_id))
        # 暂存地区关联关系
        for r in offer.regions:
            rid = get_region_id(conn, r)
            offer_region_rows.append((offer.title, cat_id, rid))
    
    with conn.cursor() as cur:
        # 批量插入所有Offer
        execute_values(cur, """
            INSERT INTO Offers(title, category) VALUES %s
            ON CONFLICT (title, category) DO NOTHING
        """, offer_rows)
        
        # 批量插入关联关系，用子查询匹配Offer ID
        execute_values(cur, """
            INSERT INTO OfferRegions(offer, region)
            SELECT o.id, t.region FROM (VALUES %s) AS t(title, cat, region)
            JOIN Offers o ON o.title = t.title AND o.category = t.cat
            ON CONFLICT (offer, region) DO NOTHING
        """, offer_region_rows)
    conn.commit()

优化效果对比 #

内容的提问来源于stack exchange，提问作者Jakub Biały