# -*- coding:utf-8 -*- """PG 连接工厂 + 薄查询封装(pg8000)。 ds_ref(形如 postgresql/prd-hobby)→ 解析项目同级 ../datasource/{ds_ref}.ini → pg8000 活连接。ds_ref → 凭据 的解析复用 datax 的 DataSourceFactory。 """ import os import re from dw_base.datax.datasources.data_source_factory import DataSourceFactory from dw_base.datax.datax_constants import DS_POSTGRE_SQL_JDBC_URL # dw_base/io/db/postgresql.py → 上 4 层到项目根 _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname( os.path.dirname(os.path.abspath(__file__))))) def parse_jdbc_url(jdbc_url): """从 jdbc:postgresql://host:port/database 抽 (host, port, database)。""" m = re.match(r'jdbc:postgresql://([^:/]+)(?::(\d+))?/(.+)', jdbc_url) if not m: raise ValueError('无法解析 PG jdbcUrl: ' + jdbc_url) return m.group(1), int(m.group(2) or 5432), m.group(3) def _resolve_datasource(ds_ref): """ds_ref → DataSource。 ds_ref 形如 'postgresql/prd-hobby',首段为 db_type(同父目录名)。 datasource ini 落点:项目同级 ../datasource/{ds_ref}.ini。 """ ds_type = ds_ref.split('/')[0] if ds_type != 'postgresql': raise NotImplementedError('暂只支持 postgresql 数据源,收到: ' + ds_type) ds_file_path = os.path.normpath( os.path.join(_PROJECT_ROOT, '..', 'datasource', ds_ref + '.ini')) if not os.path.isfile(ds_file_path): raise FileNotFoundError('数据源 ini 不存在: ' + ds_file_path) return DataSourceFactory.get_data_source(ds_type, ds_file_path) def resolve(ds_ref): """ds_ref → {host, port, database, username, password}。""" d = _resolve_datasource(ds_ref).parse() host, port, database = parse_jdbc_url(d[DS_POSTGRE_SQL_JDBC_URL]) return { 'host': host, 'port': port, 'database': database, 'username': d['username'], 'password': d['password'], } def connect(ds_ref): """ds_ref → pg8000 活连接。调用方负责 close。""" info = resolve(ds_ref) import pg8000.dbapi return pg8000.dbapi.connect( host=info['host'], port=info['port'], database=info['database'], user=info['username'], password=info['password'], ) def query(conn, sql, params=None): """跑 SQL 返回 fetchall。params 为 None 时不传参(避开空参数风格差异)。""" cur = conn.cursor() try: if params is None: cur.execute(sql) else: cur.execute(sql, params) return cur.fetchall() finally: cur.close()