"""PG 同步写入层:Orders / Goods / Merchant 的 upsert + 同步所需的查询。 忠实移植自 inkreach-bi-auth-mng 的 sync/lib/db.service.ts。 - Orders 主键:("orderNo", "createdAt") - Goods 主键:("orderNo", "productionNo") 关键:写入前用库中已有 createdAt 覆盖,避免同一订单因 createdAt 不同产生重复主键行。 """ from __future__ import annotations from datetime import datetime from typing import Any from psycopg2.extras import execute_values from flows.shared.mappers import GoodsUpsertInput, OrderUpsertInput ORDER_CHUNK = 1000 GOODS_CHUNK = 2000 LOOKUP_BATCH = 10000 class SyncDB: def __init__(self, block_name: str = "pg-inkreach-dmp") -> None: from prefect_sqlalchemy import SqlAlchemyConnector self._engine = SqlAlchemyConnector.load(block_name).get_engine() def _conn(self): return self._engine.raw_connection() # ---------- 查询辅助 ---------- def get_latest_ordered_at(self, system_name: str) -> datetime | None: with self._conn() as conn: with conn.cursor() as cur: cur.execute( 'SELECT MAX(COALESCE("orderedAt","createdAt")) FROM "Orders" WHERE "systemName"=%s', (system_name,), ) row = cur.fetchone() return row[0] if row and row[0] else None def get_pending_orders(self, system_name: str, exclude_statuses: list[str]) -> list[str]: with self._conn() as conn: with conn.cursor() as cur: cur.execute( 'SELECT DISTINCT "orderNo" FROM "Orders" WHERE "systemName"=%s ' 'AND "systemStatus" IS DISTINCT FROM ALL %s', (system_name, exclude_statuses), ) rows = cur.fetchall() return [r[0] for r in rows] def get_order_count_by_status(self, system_name: str, start: datetime, end: datetime) -> dict[str, int]: with self._conn() as conn: with conn.cursor() as cur: cur.execute( 'SELECT "systemStatus", COUNT(*) FROM "Orders" WHERE "systemName"=%s ' 'AND "orderedAt">=%s AND "orderedAt"<%s GROUP BY "systemStatus"', (system_name, start, end), ) rows = cur.fetchall() return {r[0]: int(r[1]) for r in rows} def get_local_order_nos(self, system_name: str, start: datetime, end: datetime) -> list[str]: with self._conn() as conn: with conn.cursor() as cur: cur.execute( 'SELECT DISTINCT "orderNo" FROM "Orders" WHERE "systemName"=%s ' 'AND "orderedAt">=%s AND "orderedAt"<%s', (system_name, start, end), ) rows = cur.fetchall() return [r[0] for r in rows] def get_local_goods_count(self, system_name: str, start: datetime, end: datetime) -> int: with self._conn() as conn: with conn.cursor() as cur: cur.execute( 'SELECT COUNT(*) FROM "Goods" g JOIN "Orders" o ' 'ON g."orderNo"=o."orderNo" AND g."createdAt"=o."createdAt" ' 'WHERE o."systemName"=%s AND o."orderedAt">=%s AND o."orderedAt"<%s', (system_name, start, end), ) row = cur.fetchone() return int(row[0]) if row else 0 def get_orders_with_100_prefix(self, system_name: str, statuses: list[str]) -> list[tuple[str, datetime]]: with self._conn() as conn: with conn.cursor() as cur: cur.execute( 'SELECT DISTINCT o."orderNo", o."createdAt" FROM "Orders" o ' 'JOIN "Goods" g ON g."orderNo"=o."orderNo" AND g."createdAt"=o."createdAt" ' 'WHERE o."systemName"=%s AND o."systemStatus"=ANY(%s) AND g."productionNo" LIKE \'100%\'', (system_name, statuses), ) rows = cur.fetchall() return [(r[0], r[1]) for r in rows] # ---------- 主 upsert ---------- def sync_orders_and_goods( self, orders: list[OrderUpsertInput], goods: list[GoodsUpsertInput], delete_order_nos: list[str] | None = None, ) -> dict[str, int]: result = { "ordersCreated": 0, "ordersUpdated": 0, "goodsCreated": 0, "goodsUpdated": 0, "goodsRemoved": 0, "ordersDeleted": 0, } if not orders and not goods and not delete_order_nos: return result conn = self._conn() try: system_name = orders[0].systemName if orders else None all_order_nos = list({o.orderNo for o in orders} | {g.orderNo for g in goods}) # 1. 解析已有 createdAt,覆盖到 orders(保证幂等 upsert) existing_created: dict[str, datetime] = {} for i in range(0, len(all_order_nos), LOOKUP_BATCH): chunk = all_order_nos[i : i + LOOKUP_BATCH] with conn.cursor() as cur: cur.execute( 'SELECT "orderNo","createdAt" FROM "Orders" WHERE "orderNo"=ANY(%s)' + (f' AND "systemName"=%s' if system_name else ""), (chunk, system_name) if system_name else (chunk,), ) for r in cur.fetchall(): existing_created[r[0]] = r[1] for o in orders: if o.orderNo in existing_created: o.createdAt = existing_created[o.orderNo] # 2. 去重 orders / goods seen_o: set[str] = set() orders_dedup: list[OrderUpsertInput] = [] for o in orders: key = f'{o.orderNo}|{o.createdAt.isoformat()}' if key in seen_o: continue seen_o.add(key) orders_dedup.append(o) resolved_created: dict[str, datetime] = {} for o in orders_dedup: resolved_created.setdefault(o.orderNo, o.createdAt) seen_g: set[str] = set() goods_dedup: list[GoodsUpsertInput] = [] for g in goods: if g.orderNo in resolved_created: g.createdAt = resolved_created[g.orderNo] key = f'{g.orderNo}|{g.productionNo}' if key in seen_g: continue seen_g.add(key) goods_dedup.append(g) goods_by_order: dict[str, list[GoodsUpsertInput]] = {} for g in goods_dedup: goods_by_order.setdefault(g.orderNo, []).append(g) # 3. 分批 upsert orders + 关联 goods for i in range(0, len(orders_dedup), ORDER_CHUNK): chunk = orders_dedup[i : i + ORDER_CHUNK] related: list[GoodsUpsertInput] = [] for o in chunk: related.extend(goods_by_order.get(o.orderNo, [])) stats = self._upsert_order_batch(conn, chunk, related, system_name) for k, v in stats.items(): result[k] += v # 4. 剩余 goods(其 order 不在本次 orders 列表) orders_set = {o.orderNo for o in orders_dedup} remaining = [g for g in goods_dedup if g.orderNo not in orders_set] if remaining: c, u = self._upsert_goods_only(conn, remaining, system_name) result["goodsCreated"] += c result["goodsUpdated"] += u # 5. 删除订单 if delete_order_nos: with conn.cursor() as cur: cur.execute( 'DELETE FROM "Orders" WHERE "orderNo"=ANY(%s)' + (f' AND "systemName"=%s' if system_name else ""), (delete_order_nos, system_name) if system_name else (delete_order_nos,), ) result["ordersDeleted"] = cur.rowcount # 6. 商家占位 self._ensure_merchants(conn, orders_dedup) conn.commit() except Exception: conn.rollback() raise finally: conn.close() return result def _upsert_order_batch( self, conn, orders: list[OrderUpsertInput], goods: list[GoodsUpsertInput], system_name: str | None ) -> dict[str, int]: stats = {"ordersCreated": 0, "ordersUpdated": 0, "goodsCreated": 0, "goodsUpdated": 0, "goodsRemoved": 0} with conn.cursor() as cur: # 统计已存在 → created/updated pairs = [(o.orderNo, o.createdAt) for o in orders] cur.execute( "SELECT t.order_no, t.created_at FROM unnest(%s::text[], %s::timestamp[]) AS t(order_no,created_at) " "JOIN \"Orders\" o ON o.\"orderNo\"=t.order_no AND o.\"createdAt\"=t.created_at", _unnest_pairs(pairs), ) existing_set = {(r[0], r[1]) for r in cur.fetchall()} updated = sum(1 for o in orders if (o.orderNo, o.createdAt) in existing_set) stats["ordersUpdated"] = updated stats["ordersCreated"] = len(orders) - updated execute_values( cur, '''INSERT INTO "Orders" ( "orderNo","outOrderNo","courierSource","courierName","courierNo","merchantId", "countryRegion","factoryName","systemName","systemStatus","platformName","platformStatus", "createdAt","orderedAt","confirmedAt","completedAt","shipDeadline" ) VALUES %s ON CONFLICT ("orderNo","createdAt") DO UPDATE SET "outOrderNo"=EXCLUDED."outOrderNo","courierSource"=EXCLUDED."courierSource", "courierName"=EXCLUDED."courierName","courierNo"=EXCLUDED."courierNo", "merchantId"=EXCLUDED."merchantId","countryRegion"=EXCLUDED."countryRegion", "factoryName"=EXCLUDED."factoryName","systemName"=EXCLUDED."systemName", "systemStatus"=EXCLUDED."systemStatus","platformName"=EXCLUDED."platformName", "platformStatus"=EXCLUDED."platformStatus","orderedAt"=EXCLUDED."orderedAt", "confirmedAt"=EXCLUDED."confirmedAt","completedAt"=EXCLUDED."completedAt", "shipDeadline"=EXCLUDED."shipDeadline"''', [( o.orderNo, o.outOrderNo, o.courierSource, o.courierName, o.courierNo, o.merchantId, o.countryRegion, o.factoryName, o.systemName, o.systemStatus, o.platformName, o.platformStatus, o.createdAt, o.orderedAt, o.confirmedAt, o.completedAt, o.shipDeadline, ) for o in orders], page_size=ORDER_CHUNK, ) # 关联 goods:先删本次不再出现的,再 upsert if goods: order_nos = list({g.orderNo for g in goods}) prod_nos = [g.productionNo for g in goods] cur.execute( '''DELETE FROM "Goods" WHERE "orderNo"=ANY(%s) AND "productionNo" <> ALL(%s) AND ("orderNo","createdAt") IN ( SELECT "orderNo","createdAt" FROM "Orders" WHERE "orderNo"=ANY(%s)''' + (f' AND "systemName"=%s' if system_name else "") + ')', (order_nos, prod_nos, order_nos, system_name) if system_name else (order_nos, prod_nos, order_nos), ) stats["goodsRemoved"] += cur.rowcount g_pairs = [(g.orderNo, g.productionNo) for g in goods] cur.execute( "SELECT t.order_no,t.production_no FROM unnest(%s::text[], %s::text[]) AS t(order_no,production_no) " "JOIN \"Goods\" g ON g.\"orderNo\"=t.order_no AND g.\"productionNo\"=t.production_no", _unnest_pairs(g_pairs), ) g_existing = {(r[0], r[1]) for r in cur.fetchall()} g_updated = sum(1 for g in goods if (g.orderNo, g.productionNo) in g_existing) stats["goodsUpdated"] = g_updated stats["goodsCreated"] = len(goods) - g_updated execute_values( cur, '''INSERT INTO "Goods" ( "orderNo","productionNo","createdAt","outOrderNo","itemName","itemColor", "itemSize","itemStyleNo","quantity","status" ) VALUES %s ON CONFLICT ("orderNo","productionNo") DO UPDATE SET "createdAt"=EXCLUDED."createdAt","outOrderNo"=EXCLUDED."outOrderNo", "itemName"=EXCLUDED."itemName","itemColor"=EXCLUDED."itemColor", "itemSize"=EXCLUDED."itemSize","itemStyleNo"=EXCLUDED."itemStyleNo", "quantity"=EXCLUDED."quantity","status"=EXCLUDED."status"''', [( g.orderNo, g.productionNo, g.createdAt, g.outOrderNo, g.itemName, g.itemColor, g.itemSize, g.itemStyleNo, g.quantity, g.status, ) for g in goods], page_size=GOODS_CHUNK, ) return stats def _upsert_goods_only( self, conn, goods: list[GoodsUpsertInput], system_name: str | None ) -> tuple[int, int]: # 用库中实际 createdAt 覆盖 order_nos = list({g.orderNo for g in goods}) actual: dict[str, datetime] = {} with conn.cursor() as cur: cur.execute( 'SELECT "orderNo","createdAt" FROM "Orders" WHERE "orderNo"=ANY(%s)' + (f' AND "systemName"=%s' if system_name else ""), (order_nos, system_name) if system_name else (order_nos,), ) for r in cur.fetchall(): actual[r[0]] = r[1] for g in goods: if g.orderNo in actual: g.createdAt = actual[g.orderNo] created = updated = 0 for i in range(0, len(goods), GOODS_CHUNK): chunk = goods[i : i + GOODS_CHUNK] with conn.cursor() as cur: c_order = list({g.orderNo for g in chunk}) c_prod = [g.productionNo for g in chunk] cur.execute( '''DELETE FROM "Goods" WHERE "orderNo"=ANY(%s) AND "productionNo" <> ALL(%s)''' + (f''' AND ("orderNo","createdAt") IN (SELECT "orderNo","createdAt" FROM "Orders" WHERE "orderNo"=ANY(%s) AND "systemName"=%s)''' if system_name else ""), (c_order, c_prod, c_order, system_name) if system_name else (c_order, c_prod), ) g_pairs = [(g.orderNo, g.productionNo) for g in chunk] cur.execute( "SELECT t.order_no,t.production_no FROM unnest(%s::text[], %s::text[]) AS t(order_no,production_no) " "JOIN \"Goods\" g ON g.\"orderNo\"=t.order_no AND g.\"productionNo\"=t.production_no", _unnest_pairs(g_pairs), ) g_existing = {(r[0], r[1]) for r in cur.fetchall()} u = sum(1 for g in chunk if (g.orderNo, g.productionNo) in g_existing) updated += u created += len(chunk) - u execute_values( cur, '''INSERT INTO "Goods" ( "orderNo","productionNo","createdAt","outOrderNo","itemName","itemColor", "itemSize","itemStyleNo","quantity","status" ) VALUES %s ON CONFLICT ("orderNo","productionNo") DO UPDATE SET "createdAt"=EXCLUDED."createdAt","outOrderNo"=EXCLUDED."outOrderNo", "itemName"=EXCLUDED."itemName","itemColor"=EXCLUDED."itemColor", "itemSize"=EXCLUDED."itemSize","itemStyleNo"=EXCLUDED."itemStyleNo", "quantity"=EXCLUDED."quantity","status"=EXCLUDED."status"''', [( g.orderNo, g.productionNo, g.createdAt, g.outOrderNo, g.itemName, g.itemColor, g.itemSize, g.itemStyleNo, g.quantity, g.status, ) for g in chunk], page_size=GOODS_CHUNK, ) return created, updated def _ensure_merchants(self, conn, orders: list[OrderUpsertInput]) -> None: ids = list({o.merchantId for o in orders if o.merchantId and o.merchantId.strip()}) if not ids: return with conn.cursor() as cur: execute_values( cur, '''INSERT INTO "Merchant" ("merchantId","merchantGroup","responsibleName","responsibleId","tags") VALUES %s ON CONFLICT ("merchantId") DO NOTHING''', [(mid, "", None, None, "") for mid in ids], ) def _unnest_pairs(pairs: list[tuple]) -> tuple[list, list]: """构造 psycopg2 unnest 双数组参数:返回 (col_a_list, col_b_list)。 配合 SQL:unnest(%s::text[], %s::timestamp[]) 或 unnest(%s::text[], %s::text[])。 """ return [p[0] for p in pairs], [p[1] for p in pairs]