在准备环境前提交次全部更改。
This commit is contained in:
0
apps/etl/connectors/feiqiu/database/__init__.py
Normal file
0
apps/etl/connectors/feiqiu/database/__init__.py
Normal file
112
apps/etl/connectors/feiqiu/database/base.py
Normal file
112
apps/etl/connectors/feiqiu/database/base.py
Normal file
@@ -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
|
||||
80
apps/etl/connectors/feiqiu/database/connection.py
Normal file
80
apps/etl/connectors/feiqiu/database/connection.py
Normal file
@@ -0,0 +1,80 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""数据库连接管理器（限制最大连接超时时间）。"""
|
||||
|
||||
import psycopg2
|
||||
import psycopg2.extras
|
||||
|
||||
|
||||
class DatabaseConnection:
|
||||
"""封装 psycopg2 连接，支持会话参数和超时保护。"""
|
||||
|
||||
def __init__(self, dsn: str, session: dict = None, connect_timeout: int = None):
|
||||
self._dsn = dsn
|
||||
self._session = session or {}
|
||||
self._connect_timeout = connect_timeout
|
||||
self.conn = self._open_connection()
|
||||
|
||||
def _open_connection(self):
|
||||
"""创建并初始化连接（包含会话参数）。"""
|
||||
timeout_val = self._connect_timeout if self._connect_timeout is not None else 5
|
||||
# 生产环境要求：数据库连接超时不得超过 20 秒。
|
||||
timeout_val = max(1, min(int(timeout_val), 20))
|
||||
|
||||
conn = psycopg2.connect(self._dsn, connect_timeout=timeout_val)
|
||||
conn.autocommit = False
|
||||
|
||||
# 会话参数（时区、语句超时等）
|
||||
if self._session:
|
||||
with conn.cursor() as c:
|
||||
if self._session.get("timezone"):
|
||||
c.execute("SET TIME ZONE %s", (self._session["timezone"],))
|
||||
if self._session.get("statement_timeout_ms") is not None:
|
||||
c.execute(
|
||||
"SET statement_timeout = %s",
|
||||
(int(self._session["statement_timeout_ms"]),),
|
||||
)
|
||||
if self._session.get("lock_timeout_ms") is not None:
|
||||
c.execute(
|
||||
"SET lock_timeout = %s", (int(self._session["lock_timeout_ms"]),)
|
||||
)
|
||||
if self._session.get("idle_in_tx_timeout_ms") is not None:
|
||||
c.execute(
|
||||
"SET idle_in_transaction_session_timeout = %s",
|
||||
(int(self._session["idle_in_tx_timeout_ms"]),),
|
||||
)
|
||||
return conn
|
||||
|
||||
def query(self, sql: str, args=None):
|
||||
"""Execute a query and fetch all rows."""
|
||||
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as c:
|
||||
c.execute(sql, args)
|
||||
return c.fetchall()
|
||||
|
||||
def execute(self, sql: str, args=None):
|
||||
"""Execute a SQL statement without returning rows."""
|
||||
with self.conn.cursor() as c:
|
||||
c.execute(sql, args)
|
||||
|
||||
def commit(self):
|
||||
"""Commit current transaction."""
|
||||
self.conn.commit()
|
||||
|
||||
def rollback(self):
|
||||
"""Rollback current transaction."""
|
||||
self.conn.rollback()
|
||||
|
||||
def close(self):
|
||||
"""Safely close the connection."""
|
||||
try:
|
||||
self.conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def ensure_open(self) -> bool:
|
||||
"""确保连接可用，若已关闭则尝试重连。"""
|
||||
try:
|
||||
if getattr(self.conn, "closed", 0):
|
||||
self.conn = self._open_connection()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
107
apps/etl/connectors/feiqiu/database/operations.py
Normal file
107
apps/etl/connectors/feiqiu/database/operations.py
Normal file
@@ -0,0 +1,107 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""数据库批量操作"""
|
||||
import psycopg2.extras
|
||||
import re
|
||||
|
||||
class DatabaseOperations:
|
||||
"""数据库批量操作封装"""
|
||||
|
||||
def __init__(self, connection):
|
||||
self._connection = connection
|
||||
self.conn = connection.conn
|
||||
|
||||
def batch_execute(self, sql: str, rows: list, page_size: int = 1000):
|
||||
"""批量执行SQL"""
|
||||
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,
|
||||
page_size: int = 1000) -> tuple:
|
||||
"""批量UPSERT并返回插入/更新计数"""
|
||||
if not rows:
|
||||
return (0, 0)
|
||||
|
||||
use_returning = "RETURNING" in sql.upper()
|
||||
|
||||
# 不带 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)
|
||||
|
||||
# 尝试向量化执行（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:
|
||||
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)
|
||||
except Exception:
|
||||
# 向量化失败后，事务通常处于 aborted 状态，需要先 rollback 才能继续执行。
|
||||
vectorized_failed = True
|
||||
|
||||
if vectorized_failed:
|
||||
try:
|
||||
self.conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 回退：逐行执行
|
||||
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)
|
||||
|
||||
@staticmethod
|
||||
def _is_inserted(rec) -> bool:
|
||||
"""判断是否为插入操作"""
|
||||
if rec is None:
|
||||
return False
|
||||
if isinstance(rec, tuple):
|
||||
return bool(rec[0])
|
||||
if isinstance(rec, dict):
|
||||
return bool(rec.get("inserted"))
|
||||
return False
|
||||
|
||||
# --- 透传辅助方法 -------------------------------------------------
|
||||
def commit(self):
|
||||
"""提交事务（委托给底层连接）"""
|
||||
self._connection.commit()
|
||||
|
||||
def rollback(self):
|
||||
"""回滚事务（委托给底层连接）"""
|
||||
self._connection.rollback()
|
||||
|
||||
def query(self, sql: str, args=None):
|
||||
"""执行查询并返回结果"""
|
||||
return self._connection.query(sql, args)
|
||||
|
||||
def execute(self, sql: str, args=None):
|
||||
"""执行任意 SQL"""
|
||||
self._connection.execute(sql, args)
|
||||
|
||||
def cursor(self):
|
||||
"""暴露原生 cursor，供特殊操作使用"""
|
||||
return self.conn.cursor()
|
||||
Reference in New Issue
Block a user