代码迁移

etl_billiards/database/base.py

@@ -0,0 +1,112 @@
+# -*- coding: utf-8 -*-
+"""
+数据库操作（批量、RETURNING支持）
+"""
+import re
+from typing import List, Dict, Tuple
+import psycopg2.extras
+from .connection import DatabaseConnection
+
+
+class DatabaseOperations(DatabaseConnection):
+    """扩展数据库操作（包含批量upsert和returning支持）"""
+    
+    def batch_execute(self, sql: str, rows: List[Dict], page_size: int = 1000):
+        """批量执行SQL（不带RETURNING）"""
+        if not rows:
+            return
+        with self.conn.cursor() as c:
+            psycopg2.extras.execute_batch(c, sql, rows, page_size=page_size)
+    
+    def batch_upsert_with_returning(self, sql: str, rows: List[Dict], page_size: int = 1000) -> Tuple[int, int]:
+        """
+        批量 UPSERT 并统计插入/更新数
+        
+        Args:
+            sql: 包含RETURNING子句的SQL
+            rows: 数据行列表
+            page_size: 批次大小
+            
+        Returns:
+            (inserted_count, updated_count) 元组
+        """
+        if not rows:
+            return (0, 0)
+        
+        use_returning = "RETURNING" in sql.upper()
+        
+        with self.conn.cursor() as c:
+            if not use_returning:
+                psycopg2.extras.execute_batch(c, sql, rows, page_size=page_size)
+                return (0, 0)
+            
+            # 优先尝试向量化执行
+            try:
+                inserted, updated = self._execute_with_returning_vectorized(c, sql, rows, page_size)
+                return (inserted, updated)
+            except Exception:
+                # 回退到逐行执行
+                return self._execute_with_returning_row_by_row(c, sql, rows)
+    
+    def _execute_with_returning_vectorized(self, cursor, sql: str, rows: List[Dict], page_size: int) -> Tuple[int, int]:
+        """向量化执行（使用execute_values）"""
+        # 解析VALUES子句
+        m = re.search(r"VALUES\s*\((.*?)\)", sql, flags=re.IGNORECASE | re.DOTALL)
+        if not m:
+            raise ValueError("Cannot parse VALUES clause")
+        
+        tpl = "(" + m.group(1) + ")"
+        base_sql = sql[:m.start()] + "VALUES %s" + sql[m.end():]
+        
+        ret = psycopg2.extras.execute_values(
+            cursor, base_sql, rows, template=tpl, page_size=page_size, fetch=True
+        )
+        
+        if not ret:
+            return (0, 0)
+        
+        inserted = 0
+        for rec in ret:
+            flag = self._extract_inserted_flag(rec)
+            if flag:
+                inserted += 1
+        
+        return (inserted, len(ret) - inserted)
+    
+    def _execute_with_returning_row_by_row(self, cursor, sql: str, rows: List[Dict]) -> Tuple[int, int]:
+        """逐行执行（回退方案）"""
+        inserted = 0
+        updated = 0
+        
+        for r in rows:
+            cursor.execute(sql, r)
+            try:
+                rec = cursor.fetchone()
+            except Exception:
+                rec = None
+            
+            flag = self._extract_inserted_flag(rec) if rec else None
+            
+            if flag:
+                inserted += 1
+            else:
+                updated += 1
+        
+        return (inserted, updated)
+    
+    @staticmethod
+    def _extract_inserted_flag(rec) -> bool:
+        """从返回记录中提取inserted标志"""
+        if isinstance(rec, tuple):
+            return bool(rec[0])
+        elif isinstance(rec, dict):
+            return bool(rec.get("inserted"))
+        else:
+            try:
+                return bool(rec["inserted"])
+            except Exception:
+                return False
+
+
+# 为了向后兼容，提供Pg别名
+Pg = DatabaseOperations