poyee-data-warehouse/dw_base/io/db/postgresql.py

# -*- 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()