合并

etl_billiards/database/operations.py

@@ -25,46 +25,54 @@ class DatabaseOperations:
         
         use_returning = "RETURNING" in sql.upper()
         
-        with self.conn.cursor() as c:
-            if not use_returning:
+        # 不带 RETURNING：直接批量执行即可
+        if not use_returning:
+            with self.conn.cursor() as c:
                 psycopg2.extras.execute_batch(c, sql, rows, page_size=page_size)
-                return (0, 0)
-            
-            # 尝试向量化执行
+            return (0, 0)
+
+        # 尝试向量化执行（execute_values + fetch returning）
+        vectorized_failed = False
+        m = re.search(r"VALUES\s*\((.*?)\)", sql, flags=re.IGNORECASE | re.DOTALL)
+        if m:
+            tpl = "(" + m.group(1) + ")"
+            base_sql = sql[:m.start()] + "VALUES %s" + sql[m.end():]
             try:
-                m = re.search(r"VALUES\s*\((.*?)\)", sql, flags=re.IGNORECASE | re.DOTALL)
-                if m:
-                    tpl = "(" + m.group(1) + ")"
-                    base_sql = sql[:m.start()] + "VALUES %s" + sql[m.end():]
-                    
+                with self.conn.cursor() as c:
                     ret = psycopg2.extras.execute_values(
                         c, base_sql, rows, template=tpl, page_size=page_size, fetch=True
                     )
-                    
-                    if not ret:
-                        return (0, 0)
-                    
-                    inserted = sum(1 for rec in ret if self._is_inserted(rec))
-                    return (inserted, len(ret) - inserted)
+                if not ret:
+                    return (0, 0)
+                inserted = sum(1 for rec in ret if self._is_inserted(rec))
+                return (inserted, len(ret) - inserted)
+            except Exception:
+                # 向量化失败后，事务通常处于 aborted 状态，需要先 rollback 才能继续执行。
+                vectorized_failed = True
+
+        if vectorized_failed:
+            try:
+                self.conn.rollback()
             except Exception:
                 pass
-            
-            # 回退：逐行执行
-            inserted = 0
-            updated = 0
+
+        # 回退：逐行执行
+        inserted = 0
+        updated = 0
+        with self.conn.cursor() as c:
             for r in rows:
                 c.execute(sql, r)
                 try:
                     rec = c.fetchone()
                 except Exception:
                     rec = None
-                
+
                 if self._is_inserted(rec):
                     inserted += 1
                 else:
                     updated += 1
-            
-            return (inserted, updated)
+
+        return (inserted, updated)
     
     @staticmethod
     def _is_inserted(rec) -> bool: