53 lines
1.3 KiB
Python
53 lines
1.3 KiB
Python
|
"""连接与读写 IO。PG 用 psycopg(3)，MySQL 用 pymysql。桥不 import 基座/平台代码。"""
|
|||
|
|
from contextlib import contextmanager
|
|||
|
|
|
|||
|
|
import pandas as pd
|
|||
|
|
import psycopg
|
|||
|
|
import pymysql
|
|||
|
|
|
|||
|
|
import config
|
|||
|
|
|
|||
|
|
|
|||
|
|
@contextmanager
|
|||
|
|
def akg_pg_conn():
|
|||
|
|
c = config.akg_pg()
|
|||
|
|
conn = psycopg.connect(host=c.host, port=c.port, user=c.user,
|
|||
|
|
password=c.password, dbname=c.db)
|
|||
|
|
try:
|
|||
|
|
yield conn
|
|||
|
|
finally:
|
|||
|
|
conn.close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _mysql(c: "config.Conn"):
|
|||
|
|
return pymysql.connect(host=c.host, port=c.port, user=c.user, password=c.password,
|
|||
|
|
database=c.db, charset="utf8mb4", read_timeout=180)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@contextmanager
|
|||
|
|
def _mysql_cm(which: str):
|
|||
|
|
c = {"heat": config.heat_mysql, "factor": config.factor_mysql,
|
|||
|
|
"price": config.price_mysql}[which]()
|
|||
|
|
conn = _mysql(c)
|
|||
|
|
try:
|
|||
|
|
yield conn
|
|||
|
|
finally:
|
|||
|
|
conn.close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def read_pg(sql: str, params=None) -> pd.DataFrame:
|
|||
|
|
with akg_pg_conn() as conn:
|
|||
|
|
return pd.read_sql(sql, conn, params=params)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def read_mysql(which: str, sql: str, params=None) -> pd.DataFrame:
|
|||
|
|
with _mysql_cm(which) as conn:
|
|||
|
|
return pd.read_sql(sql, conn, params=params)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@contextmanager
|
|||
|
|
def factor_conn():
|
|||
|
|
"""写因子库用（需要游标提交）。"""
|
|||
|
|
with _mysql_cm("factor") as conn:
|
|||
|
|
yield conn
|