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
|