datus-postgresql 0.1.2__tar.gz → 0.1.5__tar.gz

{datus_postgresql-0.1.2 → datus_postgresql-0.1.5}/.gitignore

@@ -138,3 +138,6 @@ Thumbs.db
 
 
 .omc
+
+uv.toml
+*/uv.toml

{datus_postgresql-0.1.2 → datus_postgresql-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datus-postgresql
-Version: 0.1.2
+Version: 0.1.5
 Summary: PostgreSQL database adapter for Datus
 Project-URL: Homepage, https://github.com/Datus-ai/datus-db-adapters
 Project-URL: Repository, https://github.com/Datus-ai/datus-db-adapters
@@ -14,8 +14,8 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Requires-Python: >=3.12
-Requires-Dist: datus-db-core>=0.1.0
-Requires-Dist: datus-sqlalchemy>=0.1.2
+Requires-Dist: datus-db-core>=0.1.3
+Requires-Dist: datus-sqlalchemy>=0.1.6
 Requires-Dist: psycopg2-binary>=2.9.11
 Requires-Dist: pydantic>=2.0.0
 Description-Content-Type: text/markdown

{datus_postgresql-0.1.2 → datus_postgresql-0.1.5}/datus_postgresql/config.py

@@ -15,10 +15,15 @@ class PostgreSQLConfig(BaseModel):
     host: str = Field(default="127.0.0.1", description="PostgreSQL server host")
     port: int = Field(default=5432, description="PostgreSQL server port")
     username: str = Field(..., description="PostgreSQL username")
-    password: str = Field(default="", description="PostgreSQL password", json_schema_extra={"input_type": "password"})
+    password: str = Field(
+        default="",
+        description="PostgreSQL password",
+        json_schema_extra={"input_type": "password"},
+    )
     database: Optional[str] = Field(default=None, description="Default database name")
     schema_name: Optional[str] = Field(default="public", alias="schema", description="Default schema name")
     sslmode: str = Field(
-        default="prefer", description="SSL mode (disable, allow, prefer, require, verify-ca, verify-full)"
+        default="prefer",
+        description="SSL mode (disable, allow, prefer, require, verify-ca, verify-full)",
     )
     timeout_seconds: int = Field(default=30, description="Connection timeout in seconds")

{datus_postgresql-0.1.2 → datus_postgresql-0.1.5}/datus_postgresql/connector.py

@@ -2,12 +2,22 @@
 # Licensed under the Apache License, Version 2.0.
 # See http://www.apache.org/licenses/LICENSE-2.0 for details.
 
+from collections import OrderedDict
 from typing import Any, Dict, List, Optional, Set, Union, override
 from urllib.parse import quote_plus
 
-from datus_db_core import TABLE_TYPE, DatusDbException, ErrorCode, get_logger, list_to_in_str
-from datus_sqlalchemy import SQLAlchemyConnector
 from pydantic import BaseModel, Field
+from sqlalchemy import create_engine, text
+
+from datus_db_core import (
+    TABLE_TYPE,
+    DatusDbException,
+    ErrorCode,
+    MigrationTargetMixin,
+    get_logger,
+    list_to_in_str,
+)
+from datus_sqlalchemy import SQLAlchemyConnector
 
 from .config import PostgreSQLConfig
 
@@ -43,7 +53,7 @@ def _get_metadata_config(table_type: TABLE_TYPE) -> TableMetadataNames:
     return METADATA_DICT[table_type]
 
 
-class PostgreSQLConnector(SQLAlchemyConnector):
+class PostgreSQLConnector(SQLAlchemyConnector, MigrationTargetMixin):
     """PostgreSQL database connector."""
 
     def __init__(self, config: Union[PostgreSQLConfig, dict]):
@@ -59,7 +69,6 @@ class PostgreSQLConnector(SQLAlchemyConnector):
         elif not isinstance(config, PostgreSQLConfig):
             raise TypeError(f"config must be PostgreSQLConfig or dict, got {type(config)}")
 
-        self.config = config
         self.host = config.host
         self.port = config.port
         self.username = config.username
@@ -76,9 +85,18 @@ class PostgreSQLConnector(SQLAlchemyConnector):
             f"{database}?sslmode={config.sslmode}"
         )
 
-        super().__init__(connection_string, dialect="postgresql", timeout_seconds=config.timeout_seconds)
-        self.database_name = database
-        self.schema_name = config.schema_name or "public"
+        super().__init__(
+            connection_string,
+            dialect="postgresql",
+            timeout_seconds=config.timeout_seconds,
+        )
+        # Set after super().__init__() so BaseSqlConnector doesn't overwrite
+        # with a plain ConnectionConfig (which lacks sslmode, etc.)
+        self.config = config
+        self._default_database = database
+        self._default_schema = config.schema_name or "public"
+        self._engines: OrderedDict = OrderedDict()  # LRU cache: database_name -> engine
+        self._max_engines = 8
 
     # ==================== System Resources ====================
 
@@ -90,15 +108,26 @@ class PostgreSQLConnector(SQLAlchemyConnector):
     @override
     def _sys_schemas(self) -> Set[str]:
         """System schemas to filter out."""
-        return {"pg_catalog", "information_schema", "pg_toast", "pg_temp_1", "pg_toast_temp_1"}
+        return {
+            "pg_catalog",
+            "information_schema",
+            "pg_toast",
+            "pg_temp_1",
+            "pg_toast_temp_1",
+        }
 
     # ==================== Utility Methods ====================
 
-    @staticmethod
-    def _quote_identifier(identifier: str) -> str:
-        """Safely wrap identifiers with double quotes for PostgreSQL."""
-        escaped = identifier.replace('"', '""')
-        return f'"{escaped}"'
+    # quote_identifier: uses BaseSqlConnector default (ANSI double quotes)
+
+    def _build_connection_string(self, database_name: str) -> str:
+        """Build a PostgreSQL connection string for a given database."""
+        encoded_username = quote_plus(self.username) if self.username else ""
+        encoded_password = quote_plus(self.password) if self.password else ""
+        return (
+            f"postgresql+psycopg2://{encoded_username}:{encoded_password}"
+            f"@{self.host}:{self.port}/{database_name}?sslmode={self.config.sslmode}"
+        )
 
     # ==================== Metadata Retrieval ====================
 
@@ -122,15 +151,18 @@ class PostgreSQLConnector(SQLAlchemyConnector):
             List of metadata dictionaries
         """
         self.connect()
+        database_name = database_name or self.database_name
         schema_name = schema_name or self.schema_name
 
         # Get metadata configuration
         metadata_config = _get_metadata_config(table_type)
 
         if table_type == "mv":
-            # Materialized views use pg_matviews
+            # pg_matviews is scoped to the current database connection.
+            # Use a temporary connection if a different database is requested (thread-safe).
+            safe_schema = schema_name.replace("'", "''") if schema_name else ""
             if schema_name:
-                where = f"schemaname = '{schema_name}'"
+                where = f"schemaname = '{safe_schema}'"
             else:
                 where = f"{list_to_in_str('schemaname not in', list(self._sys_schemas()))}"
 
@@ -139,10 +171,13 @@ class PostgreSQLConnector(SQLAlchemyConnector):
                 FROM pg_matviews
                 WHERE {where}
             """
+            query_result = self._execute_pandas(query, database_name=database_name)
         else:
-            # Tables and views use information_schema
+            # Tables and views use information_schema (supports table_catalog filter)
+            safe_schema = schema_name.replace("'", "''") if schema_name else ""
+            safe_db = database_name.replace("'", "''") if database_name else ""
             if schema_name:
-                where = f"table_schema = '{schema_name}'"
+                where = f"table_schema = '{safe_schema}'"
             else:
                 where = f"{list_to_in_str('table_schema not in', list(self._sys_schemas()))}"
 
@@ -154,10 +189,9 @@ class PostgreSQLConnector(SQLAlchemyConnector):
             query = f"""
                 SELECT table_schema, table_name
                 FROM information_schema.{metadata_config.info_table}
-                WHERE {where} {type_filter}
+                WHERE table_catalog = '{safe_db}' AND {where} {type_filter}
             """
-
-        query_result = self._execute_pandas(query)
+            query_result = self._execute_pandas(query, database_name=database_name)
 
         # Format results
         result = []
@@ -168,7 +202,7 @@ class PostgreSQLConnector(SQLAlchemyConnector):
                 {
                     "identifier": self.identifier(schema_name=schema, table_name=tb_name),
                     "catalog_name": "",
-                    "database_name": self.database_name,
+                    "database_name": database_name,
                     "schema_name": schema,
                     "table_name": tb_name,
                     "table_type": table_type,
@@ -190,10 +224,13 @@ class PostgreSQLConnector(SQLAlchemyConnector):
         """
         full_name = self.full_name(schema_name=schema_name, table_name=table_name)
 
+        safe_schema = schema_name.replace("'", "''") if schema_name else ""
+        safe_table = table_name.replace("'", "''") if table_name else ""
+
         if object_type == "VIEW":
             # Get view definition
             sql = f"""
-                SELECT pg_get_viewdef('{schema_name}.{table_name}'::regclass, true) as definition
+                SELECT pg_get_viewdef('{safe_schema}.{safe_table}'::regclass, true) as definition
             """
             result = self._execute_pandas(sql)
             if not result.empty and result["definition"][0]:
@@ -205,7 +242,7 @@ class PostgreSQLConnector(SQLAlchemyConnector):
             sql = f"""
                 SELECT definition
                 FROM pg_matviews
-                WHERE schemaname = '{schema_name}' AND matviewname = '{table_name}'
+                WHERE schemaname = '{safe_schema}' AND matviewname = '{safe_table}'
             """
             result = self._execute_pandas(sql)
             if not result.empty and result["definition"][0]:
@@ -221,7 +258,7 @@ class PostgreSQLConnector(SQLAlchemyConnector):
             col_defs = []
             pk_cols = []
             for col in columns:
-                col_def = f"    {self._quote_identifier(col['name'])} {col['type']}"
+                col_def = f"    {self.quote_identifier(col['name'])} {col['type']}"
                 if not col.get("nullable", True):
                     col_def += " NOT NULL"
                 if col.get("default_value"):
@@ -233,7 +270,7 @@ class PostgreSQLConnector(SQLAlchemyConnector):
             ddl = f"CREATE TABLE {full_name} (\n"
             ddl += ",\n".join(col_defs)
             if pk_cols:
-                pk_names = ", ".join(self._quote_identifier(c) for c in pk_cols)
+                pk_names = ", ".join(self.quote_identifier(c) for c in pk_cols)
                 ddl += f",\n    PRIMARY KEY ({pk_names})"
             ddl += "\n);"
             return ddl
@@ -307,7 +344,11 @@ class PostgreSQLConnector(SQLAlchemyConnector):
 
     @override
     def get_tables_with_ddl(
-        self, catalog_name: str = "", database_name: str = "", schema_name: str = "", tables: Optional[List[str]] = None
+        self,
+        catalog_name: str = "",
+        database_name: str = "",
+        schema_name: str = "",
+        tables: Optional[List[str]] = None,
     ) -> List[Dict[str, str]]:
         """Get tables with DDL statements."""
         return self._get_objects_with_ddl("table", tables, catalog_name, database_name, schema_name)
@@ -321,7 +362,11 @@ class PostgreSQLConnector(SQLAlchemyConnector):
 
     @override
     def get_schema(
-        self, catalog_name: str = "", database_name: str = "", schema_name: str = "", table_name: str = ""
+        self,
+        catalog_name: str = "",
+        database_name: str = "",
+        schema_name: str = "",
+        table_name: str = "",
     ) -> List[Dict[str, Any]]:
         """
         Get table schema using INFORMATION_SCHEMA.
@@ -338,8 +383,13 @@ class PostgreSQLConnector(SQLAlchemyConnector):
         if not table_name:
             return []
 
+        database_name = database_name or self.database_name
         schema_name = schema_name or self.schema_name
 
+        safe_db = database_name.replace("'", "''") if database_name else ""
+        safe_schema = schema_name.replace("'", "''") if schema_name else ""
+        safe_table = table_name.replace("'", "''") if table_name else ""
+
         # Use INFORMATION_SCHEMA to get schema with comments
         sql = f"""
             SELECT
@@ -357,15 +407,16 @@ class PostgreSQLConnector(SQLAlchemyConnector):
                     ON tc.constraint_name = kcu.constraint_name
                     AND tc.table_schema = kcu.table_schema
                 WHERE tc.constraint_type = 'PRIMARY KEY'
-                    AND tc.table_schema = '{schema_name}'
-                    AND tc.table_name = '{table_name}'
+                    AND tc.table_schema = '{safe_schema}'
+                    AND tc.table_name = '{safe_table}'
             ) pk ON c.column_name = pk.column_name
             LEFT JOIN pg_catalog.pg_statio_all_tables st
                 ON st.schemaname = c.table_schema AND st.relname = c.table_name
             LEFT JOIN pg_catalog.pg_description pgd
                 ON pgd.objoid = st.relid AND pgd.objsubid = c.ordinal_position
-            WHERE c.table_schema = '{schema_name}'
-              AND c.table_name = '{table_name}'
+            WHERE c.table_catalog = '{safe_db}'
+              AND c.table_schema = '{safe_schema}'
+              AND c.table_name = '{safe_table}'
             ORDER BY c.ordinal_position
         """
         query_result = self._execute_pandas(sql)
@@ -403,7 +454,9 @@ class PostgreSQLConnector(SQLAlchemyConnector):
     @override
     def get_schemas(self, catalog_name: str = "", database_name: str = "", include_sys: bool = False) -> List[str]:
         """Get list of schemas in the current database."""
-        sql = "SELECT schema_name FROM information_schema.schemata"
+        database_name = database_name or self.database_name
+        safe_db = database_name.replace("'", "''") if database_name else ""
+        sql = f"SELECT schema_name FROM information_schema.schemata WHERE catalog_name = '{safe_db}'"
         result = self._execute_pandas(sql)
         schemas = result["schema_name"].tolist()
 
@@ -420,15 +473,88 @@ class PostgreSQLConnector(SQLAlchemyConnector):
         """Get schema name for SQLAlchemy Inspector."""
         return schema_name or self.schema_name
 
+    def _get_engine(self, database_name: str = ""):
+        """Get or create engine for the given database. Thread-safe.
+
+        PostgreSQL requires different connection strings per database,
+        so each database gets its own engine with connection pool.
+        Uses LRU eviction (max 8 engines) to avoid holding too many connections.
+        """
+        db = database_name or self.database_name
+        with self._engine_lock:
+            if db in self._engines:
+                self._engines.move_to_end(db)
+                return self._engines[db]
+            conn_str = self._build_connection_string(db)
+            engine = create_engine(
+                conn_str,
+                pool_size=5,
+                max_overflow=10,
+                pool_timeout=self.timeout_seconds,
+                pool_recycle=3600,
+                pool_pre_ping=True,
+            )
+            self._engines[db] = engine
+            while len(self._engines) > self._max_engines:
+                _, evicted = self._engines.popitem(last=False)
+                try:
+                    evicted.dispose()
+                except Exception as e:
+                    logger.warning(f"Error disposing evicted engine: {e}")
+            return engine
+
     @override
-    def do_switch_context(self, catalog_name: str = "", database_name: str = "", schema_name: str = ""):
-        """Switch schema context by updating self.schema_name.
+    def _conn(self, catalog_name: str = "", database_name: str = "", schema_name: str = ""):
+        """Checkout connection from the correct per-database engine. Thread-safe.
 
-        Note: All queries use explicit schema qualification via full_name(),
-        so we only need to update self.schema_name here.
+        Overrides base _conn() to avoid writing to shared self.engine.
+        Each thread gets a connection from the engine matching its database_name.
+        """
+        from contextlib import contextmanager
+
+        @contextmanager
+        def _pg_conn():
+            effective_database = database_name or self.database_name
+            effective_schema = schema_name or self.schema_name
+            effective_catalog = catalog_name or self.catalog_name
+            engine = self._get_engine(effective_database)
+            conn = engine.connect()
+            try:
+                self.do_switch_context(conn, effective_catalog, effective_database, effective_schema)
+                yield conn
+            except Exception:
+                try:
+                    conn.rollback()
+                except Exception:
+                    pass
+                raise
+            finally:
+                conn.close()
+
+        return _pg_conn()
+
+    @override
+    def close(self):
+        """Dispose all engines (per-database pool + parent engine)."""
+        for engine in self._engines.values():
+            try:
+                engine.dispose()
+            except Exception as e:
+                logger.warning(f"Error disposing engine: {e}")
+        self._engines.clear()
+        # Dispose parent engine that may have been created via connect()/_ensure_engine()
+        super().close()
+
+    @override
+    def do_switch_context(self, conn, catalog_name: str = "", database_name: str = "", schema_name: str = ""):
+        """Apply schema context to a connection.
+
+        Database switching is handled by _conn() which picks the right engine
+        based on the effective database_name.
         """
         if schema_name:
-            self.schema_name = schema_name
+            conn.execute(text(f"SET search_path TO {self.quote_identifier(schema_name)}"))
+            conn.commit()
 
     # ==================== Sample Data ====================
 
@@ -477,7 +603,7 @@ class PostgreSQLConnector(SQLAlchemyConnector):
             return result
 
         # Otherwise get metadata and query all tables
-        metadata = self._get_metadata(table_type, "", "", schema_name)
+        metadata = self._get_metadata(table_type, "", database_name, schema_name)
         for meta in metadata:
             full_name = self.full_name(schema_name=meta["schema_name"], table_name=meta["table_name"])
             sql = f"SELECT * FROM {full_name} LIMIT {top_n}"
@@ -499,28 +625,109 @@ class PostgreSQLConnector(SQLAlchemyConnector):
 
     @override
     def identifier(
-        self, catalog_name: str = "", database_name: str = "", schema_name: str = "", table_name: str = ""
+        self,
+        catalog_name: str = "",
+        database_name: str = "",
+        schema_name: str = "",
+        table_name: str = "",
     ) -> str:
         """Generate a unique identifier for a table."""
+        database_name = database_name or self.database_name
         schema_name = schema_name or self.schema_name
+        if database_name and schema_name:
+            return f"{database_name}.{schema_name}.{table_name}"
         if schema_name:
             return f"{schema_name}.{table_name}"
         return table_name
 
     @override
     def full_name(
-        self, catalog_name: str = "", database_name: str = "", schema_name: str = "", table_name: str = ""
+        self,
+        catalog_name: str = "",
+        database_name: str = "",
+        schema_name: str = "",
+        table_name: str = "",
     ) -> str:
         """Build fully-qualified table name."""
+        database_name = database_name or self.database_name
         schema_name = schema_name or self.schema_name
+        if database_name and schema_name:
+            return f"{self.quote_identifier(database_name)}.{self.quote_identifier(schema_name)}.{self.quote_identifier(table_name)}"
         if schema_name:
-            return f"{self._quote_identifier(schema_name)}.{self._quote_identifier(table_name)}"
-        return self._quote_identifier(table_name)
+            return f"{self.quote_identifier(schema_name)}.{self.quote_identifier(table_name)}"
+        return self.quote_identifier(table_name)
 
     @override
     def _reset_filter_tables(
-        self, tables: Optional[List[str]] = None, catalog_name: str = "", database_name: str = "", schema_name: str = ""
+        self,
+        tables: Optional[List[str]] = None,
+        catalog_name: str = "",
+        database_name: str = "",
+        schema_name: str = "",
     ) -> List[str]:
         """Reset filter tables with full names."""
         schema_name = schema_name or self.schema_name
-        return super()._reset_filter_tables(tables, "", "", schema_name)
+        return super()._reset_filter_tables(tables, "", database_name, schema_name)
+
+    # ==================== MigrationTargetMixin ====================
+
+    def describe_migration_capabilities(self) -> Dict[str, Any]:
+        return {
+            "supported": True,
+            "dialect_family": "postgres-like",
+            "requires": [],  # OLTP — no distribution/partition required
+            "forbids": [
+                "DUPLICATE KEY (StarRocks-only)",
+                "DISTRIBUTED BY HASH ... BUCKETS (StarRocks-only)",
+                "ENGINE = (MySQL/ClickHouse syntax)",
+            ],
+            "type_hints": {
+                "HUGEINT": "NUMERIC(38,0) (Postgres has no HUGEINT/LARGEINT)",
+                "LARGEINT": "NUMERIC(38,0)",
+                "unbounded VARCHAR": "TEXT (prefer TEXT over unbounded VARCHAR)",
+                "TIMESTAMP WITH TIME ZONE": "TIMESTAMPTZ",
+                "JSON": "JSONB (prefer for indexing)",
+                "BOOLEAN": "BOOLEAN (no TINYINT cast needed)",
+            },
+            "example_ddl": (
+                "CREATE TABLE public.t (\n"
+                "  id BIGSERIAL PRIMARY KEY,\n"
+                "  name VARCHAR(255),\n"
+                "  created_at TIMESTAMPTZ DEFAULT now()\n"
+                ")"
+            ),
+        }
+
+    def suggest_table_layout(self, columns: List[Dict[str, Any]]) -> Dict[str, Any]:
+        # Postgres is OLTP — no distribution keys or bucketing required
+        return {}
+
+    def validate_ddl(self, ddl: str) -> List[str]:
+        errors: List[str] = []
+        upper = ddl.upper()
+
+        if "DUPLICATE KEY" in upper:
+            errors.append("DUPLICATE KEY is StarRocks-only syntax; Postgres does not support it")
+        if "BUCKETS" in upper and "DISTRIBUTED BY" in upper:
+            errors.append("DISTRIBUTED BY ... BUCKETS is StarRocks syntax; Postgres does not support it")
+        if "ENGINE =" in upper or "ENGINE=" in upper:
+            errors.append("ENGINE clause is MySQL/ClickHouse syntax; not supported in Postgres")
+        if "ORDER BY" in upper and "CREATE TABLE" in upper:
+            # Rough heuristic: top-level ORDER BY inside CREATE TABLE is ClickHouse's
+            # MergeTree syntax. Postgres allows ORDER BY inside CTAS SELECT, so this
+            # check is intentionally loose (only flags when accompanied by ENGINE).
+            if "ENGINE" in upper:
+                errors.append("ORDER BY inside CREATE TABLE is ClickHouse syntax; use CREATE INDEX in Postgres")
+
+        return errors
+
+    def map_source_type(self, source_dialect: str, source_type: str) -> Optional[str]:
+        import re as _re
+
+        base = _re.sub(r"\(.*\)", "", source_type.strip().upper()).strip()
+        overrides = {
+            "HUGEINT": "NUMERIC(38,0)",
+            "LARGEINT": "NUMERIC(38,0)",
+            "DATETIME": "TIMESTAMP",
+        }
+        return overrides.get(base)

{datus_postgresql-0.1.2 → datus_postgresql-0.1.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "datus-postgresql"
-version = "0.1.2"
+version = "0.1.5"
 description = "PostgreSQL database adapter for Datus"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -18,8 +18,8 @@ classifiers = [
 ]
 
 dependencies = [
-    "datus-db-core>=0.1.0",
-    "datus-sqlalchemy>=0.1.2",
+    "datus-db-core>=0.1.3",
+    "datus-sqlalchemy>=0.1.6",
     "psycopg2-binary>=2.9.11",
     "pydantic>=2.0.0",
 ]