pgsqldatatool 1.0.0__py3-none-any.whl

pgsqldatatool/__init__.py

@@ -0,0 +1,8 @@
+def main() -> None:
+    print("Hello from pgsqldatatool!")
+
+
+# from .until_async import write_async
+from .until_async import write_pg
+from .tools import records_to_df
+from .pgsql_connection_async import PoolSingleton

pgsqldatatool/data_clean.py

@@ -0,0 +1,230 @@
+
+import re
+import logging
+from datetime import datetime
+from zoneinfo import ZoneInfo
+import polars as pl
+
+
+logger = logging.getLogger(__name__)
+
+
+# 基础清洗 （去掉列名两端空白字符、去掉数据两段空白字符、删除空行）
+def lf_basic_clean(df:pl.LazyFrame|pl.DataFrame)->pl.LazyFrame:
+    """
+    基础清洗
+        1、去掉列名中的首尾空格和首尾换行符
+        2、删除数据两端的空白字符（"-",空格、制表符、回车、换行等）
+        3、删除完全为空的行
+    Args:
+        df (pl.LazyFrame | pl.DataFrame): 要清洗的数据集。
+    Returns:
+        pl.LazyFrame: 清洗后的数据集。
+    """
+    if not isinstance(df, pl.LazyFrame):
+        df= df.lazy()
+
+    # 去掉列名中的首尾空格和首尾换行符-------------------------------------------------
+    pattern = r'^[\s\n\r]+|[\s\n\r]+$'
+    new_names = (
+        pl.Series(df.collect_schema().names())
+        .str.replace_all(pattern, "")
+        .to_list()
+    )
+    df = df.rename(dict(zip(df.collect_schema().names(), new_names)))
+    
+
+    # 删除数据两端的空白字符（"-",空格、制表符、回车、换行等）---------------------------
+    df = df.with_columns(
+        pl.col(pl.String)
+            .str.strip_chars()
+            .str.replace("-", "")
+    )
+    # 删除空白行--------------------------------------------------
+    df = df.filter(
+        pl.any_horizontal(pl.all().is_not_null())
+    )    
+
+    return df
+
+# 根据字典重命名列
+def lf_rename_cols(df:pl.LazyFrame|pl.DataFrame, rename_cols: dict)->pl.LazyFrame:
+    """
+    根据字典重命名列
+    Args:
+        df (pl.LazyFrame | pl.DataFrame): 要清洗的数据集。
+        rename_cols (dict, optional): 列名映射字典。
+    Returns:
+        df.LazyFrame: 数据集
+    """
+    if not isinstance(df, pl.LazyFrame):
+        df= df.lazy()
+
+        # 根据参数修改存在的列的列名
+    if rename_cols is not None:
+        df = df.rename({
+            k: v for k, v in rename_cols.items() if k in df.columns
+        })
+
+    return df
+
+# 移除重复行
+def lf_remove_dup_rows(df:pl.LazyFrame|pl.DataFrame)->pl.LazyFrame:
+    """
+    保留第一次出现的行,删除其余完全相同的重复行
+    Args:
+        df (pl.LazyFrame|pl.DataFrame): 数据集
+
+    Returns:
+        pl.LazyFrame: 数据集
+
+    """
+    if isinstance(df, pl.DataFrame): # 转换为懒加载
+        df = df.lazy()
+    df = df.unique(
+        subset=None,      # 所有列
+        keep="first"      # 保留第一个
+    )
+    return df
+
+# 删除指定列
+def lf_remove_cols(df, cols:list[str] | None = None)->pl.LazyFrame:
+    """
+    删除指定列
+    Args:
+        df (pl.LazyFrame|pl.DataFrame): 数据集
+        cols (list[str] | None): 需要删除的列名。
+    Returns:
+        pl.LazyFrame|pl.DataFrame: 处理后的数据集
+    """
+    if isinstance(df, pl.DataFrame): # 转换为懒加载
+        df = df.lazy()
+
+    # 根据参数删除存在的列
+    if cols is not None:
+        df = df.drop([c for c in cols if c in df.columns])
+        
+    return df
+
+# 移除数字千分号
+def lf_remove_per_mille(df, cols: "all" | list[str] | None = None):
+    """
+    移除字符串数字中的千分号
+    Args:
+        df (pl.LazyFrame|pl.DataFrame): 数据集
+        cols (all | list[str] | None): 需要将千分比转换为小数列的列名。\n
+            默认为None。当值为all时，将所有列名转换为小数,返回仍是字符串。
+            
+    Returns:
+        pl.LazyFrame: 数据集
+    """
+    if not isinstance(df, pl.LazyFrame):
+        df= df.lazy()
+
+    # 移除字符串数字中的千分号-------------------------------------------------------
+    if cols:
+        if isinstance(cols, str) and cols=="all":
+            df = df.with_columns(
+                pl.when(
+                    pl.col(pl.String).str.contains(r'^\d{1,3}(,\d{3})*(\.\d+)?$')
+                )
+                .then(pl.col(pl.String).str.replace_all(",", ""))
+                .otherwise(pl.col(pl.String))
+            )
+        elif isinstance(cols, list):
+            df = df.with_columns(
+                pl.when(
+                    pl.col(cols).str.contains(r'^\d{1,3}(,\d{3})*(\.\d+)?$')
+                )
+                .then(pl.col(cols).str.replace_all(",", ""))
+                .otherwise(pl.col(cols))
+            )
+        else:
+            raise ValueError("per_mille_cols参数错误")    
+
+# 移除数字百分号
+def lf_remove_percent(df, cols: "all" | list[str] | None = None):
+    """
+    移除字符串数字中的百分号
+    Args:
+        df (pl.LazyFrame|pl.DataFrame): 数据集
+        cols (all | list[str] | None): 需要将百分比转换为小数列的列名。\n
+            默认为None。当值为all时，将所有列名转换为小数,返回仍是字符串。
+    Returns:
+        pl.LazyFrame: 数据集
+    """
+    if not isinstance(df, pl.LazyFrame):
+        df= df.lazy()
+
+
+    if isinstance(cols, str) and cols == "all":
+        target = pl.col(pl.String)
+    elif isinstance(cols, list):
+        target = pl.col(cols)
+    else:
+        raise ValueError("percent_cols参数错误")
+
+    df = df.with_columns(
+        target                       
+        .str.strip_chars()
+        .str.replace_all(",", "")
+        .str.extract(r"^(\d+\.?\d*)%$")
+        .cast(pl.Float64)
+        .truediv(100)
+        .fill_null(target)             
+    )
+    return df
+
+# 添加时间列
+def lf_add_time(df, time_zone='Asia/Shanghai') -> pl.LazyFrame:
+    """
+    添加时间列
+    Args:
+        df (pl.LazyFrame|pl.DataFrame): 数据集
+        time_zone (str, optional): 时区。默认为'Asia/Shanghai'。
+    Returns:
+        pl.LazyFrame: 数据集
+    """
+    if not isinstance(df, pl.LazyFrame):
+        df= df.lazy()
+
+    return df.with_columns(
+        # 使用 pl.lit(datetime.now(...)) 依然是安全的，
+        # 因为整个 with_columns 表达式在 LazyFrame 下是惰性的，
+        # datetime.now() 会在最终 .collect() 时才被求值！
+        pl.lit(datetime.now(ZoneInfo(time_zone))).alias('数据写入时间')
+    )
+
+# 删除完全为空的列
+def df_drop_empty_cols(df) -> pl.LazyFrame:
+    """
+    删除完全为空的列 \n
+    这个函数不是延迟执行的 \n
+    collect() 会触发运行一次计算，获取所有列的空值状态
+    Args:
+        df (pl.LazyFrame|pl.DataFrame): 数据集
+    Returns:
+        pl.LazyFrame: 数据集
+    """
+
+    # 1. 构建一个表达式：一次性计算每一列是否全为 null
+    # pl.struct 会将所有列打包，pl.all_horizontal 会并行计算
+    is_all_null_expr = pl.all_horizontal(
+        pl.col(c).is_null().alias(c) for c in df.collect_schema().names()
+    )
+    
+    # 2. 仅触发一次收集，获取所有列的空值状态
+    # 注意：这里只计算聚合值，不扫描全表数据，速度极快
+    all_null_flags = df.select(is_all_null_expr).collect().row(0)
+    
+    # 3. 在 Python 端解析出需要删除的列名
+    cols_to_drop = [
+        col_name for col_name, is_null in zip(df.collect_schema().names(), all_null_flags)
+        if is_null
+    ]
+    
+    # 4. 延迟执行 drop 操作（drop 本身是惰性的，直到最终 collect 才生效）
+    if cols_to_drop:
+        df = df.drop(cols_to_drop)
+    
+    return df

pgsqldatatool/pgsql_connection_async.py

@@ -0,0 +1,76 @@
+
+
+import os
+import asyncpg
+import asyncio
+from contextlib import asynccontextmanager
+
+class PoolSingleton:
+    _pool = None
+    _lock = asyncio.Lock()
+
+    @classmethod
+    async def get(cls):
+        if cls._pool is None:
+            async with cls._lock:
+                if cls._pool is None or cls._pool._closed:  # double check
+
+                    pg_host = os.environ.get("PG_HOST", "localhost")
+                    pg_port = os.environ.get("PG_PORT", "5432")
+                    pg_user = os.environ.get("PG_USER", "postgres")
+                    pg_password = os.environ.get("PG_PASSWORD", "password")
+                    pg_database = os.environ.get("PG_DB", "testdb")
+                    min_size = int(os.environ.get("PG_MIN_SIZE", "5"))
+                    max_size = int(os.environ.get("PG_MAX_SIZE", "20"))
+
+                    # 写法1
+                    # cls._pool = await asyncpg.create_pool(
+                    #     # dsn="postgresql://postgres:password@localhost:5432/testdb",
+                    #     dsn = f"postgresql://{pg_user}:{pg_password}@{pg_host}:{pg_port}/{pg_database}",
+                    #     min_size=min_size,
+                    #     max_size=max_size
+                    # )
+
+                    # 独立参数写法
+                    cls._pool = await asyncpg.create_pool(
+                        user=pg_user,
+                        password=pg_password,
+                        host=pg_host,
+                        port=pg_port,
+                        database=pg_database,
+                        min_size=min_size,
+                        max_size=max_size
+                    )
+        return cls._pool
+
+    @classmethod
+    async def close(cls):
+        if cls._pool:
+            await cls._pool.close()
+            cls._pool = None
+
+    # @classmethod
+    # async def acquire(cls):
+    #     pool = await cls.get()
+    #     return pool.acquire()
+
+    @classmethod
+    @asynccontextmanager
+    async def acquire(cls):
+        pool = await cls.get()
+        async with pool.acquire() as conn:
+            yield conn
+    @classmethod
+    async def fetch(cls, query, *args):
+        async with cls.acquire() as conn:
+            return await conn.fetch(query, *args)
+
+    @classmethod
+    async def fetchrow(cls, query, *args):
+        async with cls.acquire() as conn:
+            return await conn.fetchrow(query, *args)
+
+    @classmethod
+    async def execute(cls, query, *args):
+        async with cls.acquire() as conn:
+            return await conn.execute(query, *args)

pgsqldatatool/tools.py

@@ -0,0 +1,24 @@
+import asyncpg
+import polars as pl
+
+def records_to_df(records:list[asyncpg.Record]):
+    """
+    将asyncpg
+    将asyncpg.Record对象列表转换为polars.DataFrame对象
+    理论上这个写法比 df = pl.DataFrame([dict(r) for r in records]) 更高效
+    Args:
+        records (list[asyncpg.Record]): 要转换的记录列表。
+    Returns:
+        pl.DataFrame: 转换后的DataFrame对象。
+    """
+    if not records:
+        return pl.DataFrame()
+
+    cols = list(records[0].keys())  # ✅ 转成 list
+    data = [[] for _ in range(len(cols))]
+
+    for r in records:
+        for i, val in enumerate(r):
+            data[i].append(val)
+
+    return pl.DataFrame({col: data[i] for i, col in enumerate(cols)})

pgsqldatatool/until_async.py

@@ -0,0 +1,678 @@
+"""异步 PostgreSQL 写入工具（Polars 专用）。"""
+
+import io
+import time
+from datetime import datetime
+from typing import Any, Dict, List, Literal, Optional, Union
+
+import polars as pl
+import asyncpg
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from polars.datatypes import (
+    Int8, Int16, Int32, Int64,
+    UInt8, UInt16, UInt32, UInt64,
+    Float32, Float64,
+    Boolean,
+    String,
+    Datetime,
+    Duration,
+    List,
+    Struct,
+    Decimal,
+    Date,
+    Time,
+)
+
+import logging
+
+logger = logging.getLogger("pgtool")
+
+# =========================================================
+# 公开 API
+# =========================================================
+
+
+async def write_pg(
+    df: pl.DataFrame,
+    schema: str,
+    table: str,
+    conn: Union[asyncpg.Connection, AsyncSession],
+    *,
+    create: bool = False,
+    update: bool = False,
+    add_columns: bool = False,
+    only_column: bool = False,
+    write: Literal["auto", "copy", "execute", "executemany"] = "auto",
+    primary_key: Optional[Union[str, List[str]]] = None,
+) -> int:
+    """
+    将 Polars DataFrame 异步写入 PostgreSQL 表。
+
+    支持：
+    - 自动建表（基于 Polars 原始 dtype 映射）
+    - 自动新增字段（基于 Polars 原始 dtype 映射）
+    - 字段对齐
+    - 主键校验
+    - UPSERT（ON CONFLICT）
+
+    Args:
+        df (pl.DataFrame):
+            要写入的数据。
+        schema (str):
+            数据库模式名。
+        table (str):
+            目标表名。
+        conn (Union[asyncpg.Connection, AsyncSession]):
+            异步连接对象。
+        create (bool, optional):
+            表不存在时是否自动创建，默认 False。
+        update (bool, optional):
+            是否更新已存在的数据，默认 False。
+        add_columns (bool, optional):
+            是否自动添加表中不存在的列，默认 False。
+        only_column (bool, optional):
+            是否只写入表中已存在的列，默认 False。
+        write (Literal["auto", "copy", "execute", "executemany"], optional):
+            写入模式：
+            - auto: 自动选择
+            - copy: COPY 模式（大数据量推荐）
+            - execute: 单行 INSERT（慢）
+            - executemany: 多行 INSERT
+        primary_key (Optional[Union[str, List[str]]], optional):
+            主键列名，用于自动建表和 UPSERT。
+            - None: 不设主键
+            - str: 单列主键（如 "id"、"uuid"）
+            - List[str]: 复合主键（如 ["user_id", "order_id"]）
+            仅在 create=True 且表不存在时用于建表；
+            表已存在时以数据库实际主键为准。
+
+    Returns:
+        int:
+            受影响的行数。
+
+    Raises:
+        ValueError:
+            表不存在且 create=False，
+            或字段不匹配，
+            或主键问题。
+        TypeError:
+            Polars 类型无法映射到 PG 类型。
+        RuntimeError:
+            无法解析 asyncpg 连接。
+    """
+    start_time = time.time()
+
+    if df.is_empty():
+        logger.warning("DataFrame 为空，跳过写入")
+        return 0
+
+    # 保存 Polars 原始 dtype（用于建表 / 新增字段）
+    original_schema = df.schema
+
+    actual_conn = await _get_asyncpg_connection(conn)
+
+    async with actual_conn.transaction():
+        count = await _write_data(
+            df=df,
+            original_schema=original_schema,
+            schema=schema,
+            table=table,
+            conn=actual_conn,
+            create=create,
+            update=update,
+            add_columns=add_columns,
+            only_column=only_column,
+            write=write,
+            start_time=start_time,
+            primary_key=primary_key,
+        )
+
+    logger.info("总耗时 %.2fs", time.time() - start_time)
+    return count
+
+
+# =========================================================
+# 连接适配
+# =========================================================
+
+
+async def _get_asyncpg_connection(conn: Any) -> asyncpg.Connection:
+    """
+    从支持的数据库连接类型中提取 asyncpg 连接。
+
+    Args:
+        conn (Any):
+            asyncpg.Connection 或 SQLAlchemy AsyncSession。
+
+    Returns:
+        asyncpg.Connection:
+            解析后的 asyncpg 连接。
+
+    Raises:
+        ValueError:
+            不支持的连接类型。
+    """
+    if isinstance(conn, asyncpg.Connection):
+        return conn
+
+    if isinstance(conn, AsyncSession):
+        try:
+            raw = await conn.connection()
+            return raw.driver_connection
+        except Exception as exc:
+            raise ValueError("无法从 AsyncSession 获取 asyncpg 连接") from exc
+
+    raise ValueError(f"不支持的连接类型: {type(conn)}")
+
+
+# =========================================================
+# Polars dtype → PostgreSQL 类型映射（唯一入口）
+# =========================================================
+
+
+def polars_to_pg_type(dtype) -> str:
+    """
+    将 Polars DataType 映射为 PostgreSQL 列类型。
+
+    这是整个模块中唯一的类型映射入口，
+    建表和新增字段都会调用此函数。
+
+    Args:
+        dtype:
+            Polars DataType 实例，例如 pl.Int64、pl.String 等。
+
+    Returns:
+        str:
+            PostgreSQL 列类型字符串。
+
+    Raises:
+        TypeError:
+            当 dtype 无法映射时抛出。
+    """
+    mapping = {
+        Int8: "SMALLINT",
+        Int16: "SMALLINT",
+        Int32: "INTEGER",
+        Int64: "BIGINT",
+
+        UInt8: "SMALLINT",
+        UInt16: "INTEGER",
+        UInt32: "BIGINT",
+        UInt64: "NUMERIC(20)",
+
+        Float32: "REAL",
+        Float64: "DOUBLE PRECISION",
+
+        Boolean: "BOOLEAN",
+
+        String: "TEXT",
+
+        Datetime: "TIMESTAMP",
+        Date: "DATE",
+        Time: "TIME",
+
+        Duration: "INTERVAL",
+
+        Decimal: "NUMERIC",
+    }
+
+    # List / Struct 需要特殊处理，不能直接映射为简单列
+    if dtype in (List, Struct):
+        raise TypeError(
+            f"Polars dtype {dtype} 无法映射为简单的 PG 列类型"
+        )
+
+    pg_type = mapping.get(dtype)
+    if pg_type is None:
+        raise TypeError(
+            f"不支持的 Polars dtype: {dtype}"
+        )
+
+    return pg_type
+
+
+# =========================================================
+# 主键解析
+# =========================================================
+
+
+def _resolve_primary_key(
+    df_columns: List[str],
+    primary_key: Optional[Union[str, List[str]]],
+) -> List[str]:
+    """
+    解析并校验用户指定的主键列名。
+
+    Args:
+        df_columns: DataFrame 的所有列名。
+        primary_key: 用户显式指定的主键，None 表示不设主键。
+
+    Returns:
+        List[str]: 主键列名列表（primary_key 为 None 时返回空列表）。
+
+    Raises:
+        ValueError: 指定的主键列在 DataFrame 中不存在。
+    """
+    if primary_key is None:
+        return []
+
+    if isinstance(primary_key, str):
+        pk_list = [primary_key]
+    else:
+        pk_list = list(primary_key)
+
+    missing = [c for c in pk_list if c not in df_columns]
+    if missing:
+        raise ValueError(
+            f"指定的主键列 {missing} 在 DataFrame 中不存在"
+        )
+    return pk_list
+
+
+# =========================================================
+# 核心写入逻辑
+# =========================================================
+
+
+async def _write_data(
+    df: pl.DataFrame,
+    original_schema,
+    schema: str,
+    table: str,
+    conn: asyncpg.Connection,
+    create: bool,
+    update: bool,
+    add_columns: bool,
+    only_column: bool,
+    write: str,
+    start_time: float,
+    primary_key: Optional[Union[str, List[str]]] = None,
+) -> int:
+    """
+    内部函数，实现主要写入流程。
+    """
+    # ---- 查询目标表现在已有的字段 ----
+    db_columns_info = await conn.fetch(
+        """
+        SELECT column_name, data_type
+        FROM information_schema.columns
+        WHERE table_schema = $1 AND table_name = $2
+        ORDER BY ordinal_position;
+        """,
+        schema,
+        table,
+    )
+
+    db_columns = [r["column_name"] for r in db_columns_info]
+    db_column_types: Dict[str, str] = {
+        r["column_name"]: r["data_type"] for r in db_columns_info
+    }
+    table_exists = bool(db_columns)
+
+    # =========================
+    # 建表（使用 Polars 原始 dtype）
+    # =========================
+    if not table_exists:
+        if not create:
+            raise ValueError(f"表 {schema}.{table} 不存在，且 create=False")
+
+        cols_def = []
+        for col_name, dtype in original_schema.items():
+            pg_type = polars_to_pg_type(dtype)
+            cols_def.append(f'"{col_name}" {pg_type}')
+            db_column_types[col_name] = pg_type
+
+        # 确定主键列
+        resolved_pk = _resolve_primary_key(df.columns, primary_key)
+
+        pk_def = ""
+        if resolved_pk:
+            pk_cols_sql = ", ".join(f'"{c}"' for c in resolved_pk)
+            pk_def = f',\n    PRIMARY KEY ({pk_cols_sql})'
+
+        create_sql = (
+            f'CREATE TABLE "{schema}"."{table}" (\n'
+            + ",\n".join(f"    {c}" for c in cols_def)
+            + pk_def
+            + "\n);"
+        )
+        await conn.execute(create_sql)
+
+        db_columns = list(df.columns)
+        logger.info(
+            "表 %s.%s 创建成功，字段数: %d%s",
+            schema,
+            table,
+            len(cols_def),
+            f"，主键: {resolved_pk}" if resolved_pk else "",
+        )
+
+    # =========================
+    # 字段对齐
+    # =========================
+    df_cols = list(df.columns)
+
+    if only_column:
+        common_cols = [c for c in df_cols if c in db_columns]
+        df = df.select(common_cols)
+        df_cols = common_cols
+
+    elif add_columns:
+        # 新增字段（使用 Polars 原始 dtype）
+        new_cols = [c for c in df_cols if c not in db_columns]
+        for col in new_cols:
+            pg_type = polars_to_pg_type(original_schema[col])
+            await conn.execute(
+                f'ALTER TABLE "{schema}"."{table}" ADD COLUMN "{col}" {pg_type}'
+            )
+            db_column_types[col] = pg_type
+            db_columns.append(col)
+            logger.info(
+                "新增字段 %s.%s.%s %s",
+                schema,
+                table,
+                col,
+                pg_type,
+            )
+
+    else:
+        extra = set(df_cols) - set(db_columns)
+        if extra:
+            raise ValueError(
+                f"列 {extra} 不存在于目标表，且 add_columns=False"
+            )
+
+    # =========================
+    # 主键校验
+    # =========================
+    pk_rows = await conn.fetch(
+        """
+        SELECT kcu.column_name
+        FROM information_schema.table_constraints tc
+        JOIN information_schema.key_column_usage kcu
+          ON tc.constraint_name = kcu.constraint_name
+         AND tc.table_schema = kcu.table_schema
+        WHERE tc.constraint_type = 'PRIMARY KEY'
+          AND tc.table_schema = $1
+          AND tc.table_name = $2
+        ORDER BY kcu.ordinal_position;
+        """,
+        schema,
+        table,
+    )
+    pk_cols = [r["column_name"] for r in pk_rows]
+
+    if pk_cols:
+        missing_pk = set(pk_cols) - set(df_cols)
+        if missing_pk:
+            raise ValueError(f"主键列缺失: {missing_pk}")
+
+        dup = df.group_by(pk_cols).agg(pl.len()).filter(pl.col("len") > 1)
+        if not dup.is_empty():
+            raise ValueError(f"主键重复:\n{dup.head(1)}")
+
+    # =========================
+    # 类型对齐：根据数据库实际列类型转换 DataFrame 值
+    # =========================
+    str_types = {"text", "varchar", "character varying", "char", "character", "name"}
+    for col in df_cols:
+        db_type = db_column_types.get(col, "").lower()
+        if db_type in str_types and not isinstance(df[col].dtype, pl.String):
+            df = df.with_columns(pl.col(col).cast(pl.String))
+            logger.debug("列 %s 值转换为 String（数据库类型: %s）", col, db_type)
+
+    # =========================
+    # 选择写入模式
+    # =========================
+    nrows, ncols = df.shape
+    if write == "auto":
+        # 根据数据量自动选择最优写入方式：
+        #   - < 100 行：单行 execute（避免 executemany 协议开销）
+        #   - 100 ~ 65535 单元格：批量 executemany
+        #   - > 65535 单元格：COPY 协议（大数据量最优）
+        if nrows < 100:
+            write = "execute"
+        elif nrows * ncols > 65535:
+            write = "copy"
+        else:
+            write = "executemany"
+
+    # =========================
+    # 执行写入
+    # =========================
+    if write == "copy":
+        return await _write_copy(
+            df, schema, table, conn, df_cols, pk_cols, update
+        )
+    elif write == "executemany":
+        return await _write_executemany(
+            df, schema, table, conn, df_cols, pk_cols, update
+        )
+    elif write == "execute":
+        return await _write_execute(
+            df, schema, table, conn, df_cols, pk_cols, update
+        )
+    else:
+        raise ValueError(f"不支持的写入模式: {write}")
+
+
+# =========================================================
+# COPY 模式
+# =========================================================
+
+
+async def _write_copy(
+    df: pl.DataFrame,
+    schema: str,
+    table: str,
+    conn: asyncpg.Connection,
+    df_cols: List[str],
+    pk_cols: List[str],
+    update: bool,
+) -> int:
+    """
+    使用 PostgreSQL COPY 模式写入数据（适合大数据量）。
+
+    通过临时表中转：
+    1. 创建结构与目标表相同的临时表
+    2. 使用 COPY 协议批量灌入临时表
+    3. 从临时表 INSERT INTO 目标表（支持 ON CONFLICT）
+
+    注意：临时表在事务提交/回滚时自动清理。
+    """
+    temp_table = f"_temp_{table}_{datetime.now():%Y%m%d%H%M%S%f}"
+
+    await conn.execute(
+        f'CREATE TEMP TABLE "{temp_table}" '
+        f'AS SELECT * FROM "{schema}"."{table}" WHERE 1=0'
+    )
+
+    # 流式写入 CSV 到 BytesIO，避免内存中同时存在字符串和字节副本
+    buf = io.BytesIO()
+    df.select(df_cols).write_csv(
+        file=buf,
+        separator="\t",
+        null_value=r"\N",
+        include_header=False,
+        quote_style="never",
+    )
+    buf.seek(0)
+
+    await conn.copy_to_table(
+        table_name=temp_table,
+        source=buf,
+        columns=df_cols,
+        format="csv",
+        delimiter="\t",
+        null=r"\N",
+    )
+
+    quoted_cols = [f'"{c}"' for c in df_cols]
+
+    if not pk_cols:
+        result = await conn.execute(
+            f'INSERT INTO "{schema}"."{table}" '
+            f'SELECT * FROM "{temp_table}"'
+        )
+    else:
+        if update:
+            update_set = ", ".join(
+                f'"{c}" = EXCLUDED."{c}"'
+                for c in df_cols
+                if c not in pk_cols
+            )
+            conflict_action = f"DO UPDATE SET {update_set}"
+        else:
+            conflict_action = "DO NOTHING"
+
+        result = await conn.execute(
+            f"""
+            INSERT INTO "{schema}"."{table}" ({", ".join(quoted_cols)})
+            SELECT {", ".join(quoted_cols)} FROM "{temp_table}"
+            ON CONFLICT ({", ".join(f'"{c}"' for c in pk_cols)})
+            {conflict_action}
+            """
+        )
+
+    # asyncpg execute 返回格式为 "INSERT 0 N"，取最后一个 token 即为行数
+    count = int(result.split()[-1])
+    logger.info("COPY 写入完成: %d 行", count)
+    return count
+
+
+# =========================================================
+# EXECUTEMANY 模式
+# =========================================================
+
+
+async def _write_executemany(
+    df: pl.DataFrame,
+    schema: str,
+    table: str,
+    conn: asyncpg.Connection,
+    df_cols: List[str],
+    pk_cols: List[str],
+    update: bool,
+) -> int:
+    """
+    使用 executemany 多行 INSERT 写入数据。
+
+    注意：asyncpg 的 executemany 不返回精确的受影响行数，
+    因此返回值是尝试写入的总行数（而非实际 INSERT/UPDATE 行数）。
+    """
+    # 生成行元组列表，适配 asyncpg executemany
+    records = [tuple(row) for row in df.select(df_cols).iter_rows()]
+
+    placeholders = ", ".join(f"${i+1}" for i in range(len(df_cols)))
+    quoted_cols = [f'"{c}"' for c in df_cols]
+
+    if not pk_cols:
+        if update:
+            logger.warning(
+                "表 %s.%s 无主键，update=True 无效，数据可能重复插入",
+                schema,
+                table,
+            )
+        await conn.executemany(
+            f'INSERT INTO "{schema}"."{table}" '
+            f'({", ".join(quoted_cols)}) VALUES ({placeholders})',
+            records,
+        )
+        logger.info("executemany 写入完成: %d 行", len(records))
+        return len(records)
+
+    if update:
+        update_set = ", ".join(
+            f'"{c}" = EXCLUDED."{c}"'
+            for c in df_cols
+            if c not in pk_cols
+        )
+        conflict_action = f"DO UPDATE SET {update_set}"
+    else:
+        conflict_action = "DO NOTHING"
+
+    await conn.executemany(
+        f"""
+        INSERT INTO "{schema}"."{table}" ({", ".join(quoted_cols)})
+        VALUES ({placeholders})
+        ON CONFLICT ({", ".join(f'"{c}"' for c in pk_cols)})
+        {conflict_action}
+        """,
+        records,
+    )
+
+    logger.info(
+        "executemany upsert 完成: 尝试写入 %d 行", len(records)
+    )
+    return len(records)
+
+
+# =========================================================
+# EXECUTE 模式（单行，仅适合小数据量）
+# =========================================================
+
+
+async def _write_execute(
+    df: pl.DataFrame,
+    schema: str,
+    table: str,
+    conn: asyncpg.Connection,
+    df_cols: List[str],
+    pk_cols: List[str],
+    update: bool,
+) -> int:
+    """
+    使用单行 INSERT 写入数据（仅适合极少量数据，< 100 行）。
+
+    每行一个 asyncpg.execute 调用，串行执行，效率最低。
+    大批量数据请使用 copy 或 executemany 模式。
+    """
+    count = 0
+    placeholders = ", ".join(f"${i+1}" for i in range(len(df_cols)))
+    quoted_cols = [f'"{c}"' for c in df_cols]
+
+    if not pk_cols:
+        if update:
+            logger.warning(
+                "表 %s.%s 无主键，update=True 无效，数据可能重复插入",
+                schema,
+                table,
+            )
+        sql = (
+            f'INSERT INTO "{schema}"."{table}" '
+            f'({", ".join(quoted_cols)}) VALUES ({placeholders})'
+        )
+        for row in df.iter_rows(named=True):
+            await conn.execute(sql, *(row[c] for c in df_cols))
+            count += 1
+    else:
+        if update:
+            update_set = ", ".join(
+                f'"{c}" = EXCLUDED."{c}"'
+                for c in df_cols
+                if c not in pk_cols
+            )
+            conflict_action = f"DO UPDATE SET {update_set}"
+        else:
+            conflict_action = "DO NOTHING"
+
+        sql = f"""
+            INSERT INTO "{schema}"."{table}" ({", ".join(quoted_cols)})
+            VALUES ({placeholders})
+            ON CONFLICT ({", ".join(f'"{c}"' for c in pk_cols)})
+            {conflict_action}
+        """
+
+        for row in df.iter_rows(named=True):
+            result = await conn.execute(
+                sql, *(row[c] for c in df_cols)
+            )
+            # asyncpg execute 返回 "INSERT 0 N" 格式，解析受影响行数
+            try:
+                count += int(result.split()[-1])
+            except (ValueError, IndexError):
+                count += 1
+
+    logger.info("execute 写入完成: %d 行", count)
+    return count

pgsqldatatool-1.0.0.dist-info/METADATA

@@ -0,0 +1,204 @@
+Metadata-Version: 2.3
+Name: pgsqldatatool
+Version: 1.0.0
+Summary: Add your description here
+Author: manji
+Author-email: manji <pnsm@qq.com>
+Requires-Dist: asyncpg>=0.31.0
+Requires-Dist: polars>=1.41.2
+Requires-Dist: python-dotenv>=1.2.2
+Requires-Dist: sqlalchemy>=2.0.51
+Requires-Dist: tzdata>=2026.2
+Requires-Python: >=3.14
+Description-Content-Type: text/markdown
+
+## 使用教程
+
+### 安装
+```python
+pip install pgsqldatatool
+```
+
+### date_clean（数据清洗）
+
+``` python
+# 基础清洗 （去掉列名两端空白字符、去掉数据两段空白字符、删除空行）
+lf_basic_clean(df)
+
+# 根据字典重命名列
+lf_rename_cols
+
+# 移除重复行
+lf_remove_dup_rows
+
+# 删除指定列
+lf_remove_cols
+
+# 移除数字千分号
+lf_remove_per_mille
+
+# 移除数字百分号
+lf_remove_percent
+
+# 添加时间列
+lf_add_time
+
+# 删除完全为空的列 -- 这一步不是惰性计算，可能会降低性能
+df_drop_empty_cols
+```
+示例
+
+```python
+from pgsqldatatool import data_clean as dc
+df = pl.read_excel(r"D:\manji\Downloads\判断中国域名.xlsx")
+df = dc.lf_basic_clean(df)
+df = dc.lf_remove_cols(df,["域名22"])
+df = dc.lf_remove_dup_rows(df)
+df = dc.df_drop_empty_cols(df)
+df = dc.lf_remove_dup_rows(df)
+df = df.collect()
+print(df)
+```
+
+### PoolSingleton（连接池单例模式）
+
+归还连接（Release）：把连接还给池子，让别的任务继续用。（async with 已经帮你自动做了）
+关闭连接池（Close）：彻底断开与数据库的所有连接，销毁这个池子。这通常只在整个程序/服务准备退出时才需要做。
+```python
+from pgsqldatatool import PoolSingleton
+
+# 示例1：使用异步连接池
+async def main_test_1():
+
+    # 调用方式1，
+    async with  PoolSingleton.acquire() as conn:
+        records = await conn.fetch(""" SELECT * FROM public."test20260211" """)
+        print(records)
+
+    # 销毁整个连接池
+    await PoolSingleton.close()
+
+```
+
+
+```python
+
+# 示例2：使用静态方法
+async def main_test_2():
+    # 调用方式2：使用静态方法
+    records2 = await PoolSingleton.fetch(""" SELECT * FROM public."test20260211" """)
+    print(records2)
+
+    # 调用方式3：使用静态方法
+    records3 = await PoolSingleton.fetchrow(""" SELECT * FROM public."test20260211" """)
+    print(records3)
+
+    # 调用方式4：使用静态方法
+    records4 = await PoolSingleton.execute(""" SELECT * FROM public."test20260211" """)
+    print(records4)
+    
+    
+if __name__ == "__main__":
+    asyncio.run(main_test_1())
+```
+
+
+## 创建异步数据库链接
+
+### 推荐方案：asyncpg（性能最好，原生 asyncio）
+适合：FastAPI / aiohttp / asyncio 项目
+```python
+pip install asyncpg
+```
+
+```python
+import asyncpg
+import asyncio
+
+async def main():
+    # 创建连接
+    conn = await asyncpg.connect(
+        host="localhost",
+        port=5432,
+        user="postgres",
+        password="password",
+        database="testdb"
+    )
+
+    # 查询
+    row = await conn.fetchrow("SELECT NOW()")
+    print(row)
+
+    # 参数化查询
+    rows = await conn.fetch(
+        "SELECT * FROM users WHERE age > $1",
+        18
+    )
+
+    await conn.close()
+
+asyncio.run(main())
+```
+
+### 使用连接池（生产必选 ✅）
+
+```python
+import asyncpg
+import asyncio
+
+async def get_pool():
+    return await asyncpg.create_pool(
+        dsn="postgresql://postgres:password@localhost:5432/testdb",
+        min_size=5,
+        max_size=20
+    )
+
+async def main():
+    pool = await get_pool()
+
+    async with pool.acquire() as conn:
+        result = await conn.fetch("SELECT * FROM users")
+
+    await pool.close()
+
+asyncio.run(main())
+```
+
+事务
+```python
+async with conn.transaction():
+    await conn.execute(
+        "INSERT INTO users(name) VALUES($1)",
+        "Alice"
+    )
+```
+
+
+### SQLAlchemy 2.0 + asyncpg（ORM 场景）
+适合：需要 ORM、多数据库兼容
+
+```python
+pip install sqlalchemy[asyncio] asyncpg
+```
+
+```python
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy import select
+
+engine = create_async_engine(
+    "postgresql+asyncpg://postgres:password@localhost/testdb",
+    pool_size=10,
+    max_overflow=20
+)
+
+AsyncSessionLocal = sessionmaker(
+    engine, class_=AsyncSession, expire_on_commit=False
+)
+
+async def query_users():
+    async with AsyncSessionLocal() as session:
+        result = await session.execute(select(User))
+        return result.scalars().all()
+```
+

pgsqldatatool-1.0.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+pgsqldatatool/__init__.py,sha256=Lu4cRHywCzTHUME7b-huEGa23FnCofebSgf9CrJA56Q,217
+pgsqldatatool/data_clean.py,sha256=3ZQoXsoFkZE_GyKcW1CGTDDIscIwQwdeaUHUdtlmJec,7572
+pgsqldatatool/pgsql_connection_async.py,sha256=aQxVR_vPpSFV_ozmnfzWF2-bIF_PB_fKtXK0Dmg-r5A,2609
+pgsqldatatool/tools.py,sha256=e5gqB3wsQoQ910e8JV2yfrLH301sLnh68Q7pz7VIZbk,741
+pgsqldatatool/until_async.py,sha256=pE1gmG1VCVg4jwa7ulu-PvO1OH8sNnNc2pKqrwYS1Rc,20464
+pgsqldatatool-1.0.0.dist-info/WHEEL,sha256=9sjN42GvvIkyGb9JrWAWXnA96E2dxDe0tzHzrLxUlD4,81
+pgsqldatatool-1.0.0.dist-info/entry_points.txt,sha256=58ZLpCEz7tm5xblywFKsSK6i_BeLZOsjQtQLpiOU8rY,54
+pgsqldatatool-1.0.0.dist-info/METADATA,sha256=nj0mKbKmMe_M7uyUTIbFQEZgbqyDOx3YfGTH1zASp38,4657
+pgsqldatatool-1.0.0.dist-info/RECORD,,

pgsqldatatool-1.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.10.11
+Root-Is-Purelib: true
+Tag: py3-none-any

pgsqldatatool-1.0.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+pgsqldatatool = pgsqldatatool:main
+