tai42-mcp-dynamic-postgres 0.5.0__py3-none-any.whl

tai42_mcp_dynamic_postgres/cli/main.py

@@ -0,0 +1,177 @@
+import asyncio
+import logging
+import sys
+from typing import Literal, Optional, Tuple
+
+import click
+
+from tai42_mcp_dynamic_postgres.core.app import mcp_app
+from tai42_mcp_dynamic_postgres.database.connection import close_connection_pool, get_async_connection
+from tai42_mcp_dynamic_postgres.gen.loader import load_dynamic_tools
+
+if sys.platform == "win32":
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+TransportName = Literal["stdio", "http", "sse", "streamable-http"]
+
+
+async def runner(
+    overwrite: bool,
+    readonly: bool,
+    allow_unfiltered: bool,
+    select_joined: Tuple[str, ...],
+    ignore_insert_column: Tuple[str, ...],
+    ignore_select_column: Tuple[str, ...],
+    ignore_update_column: Tuple[str, ...],
+    ignore_select_joined_column: Tuple[str, ...],
+    transport: TransportName,
+    host: str,
+    port: int,
+) -> Optional[int]:
+    if transport == "stdio" and (host != "127.0.0.1" or port != 8000):
+        raise click.BadParameter("Host and port should not be set when using 'stdio' transport.")
+
+    await load_dynamic_tools(
+        overwrite=overwrite,
+        readonly=readonly,
+        allow_unfiltered=allow_unfiltered,
+        ignore_insert_columns=list(ignore_insert_column),
+        ignore_select_columns=list(ignore_select_column),
+        ignore_update_columns=list(ignore_update_column),
+        ignore_select_joined_columns=list(ignore_select_joined_column),
+        select_joined=[group.split(",") for group in select_joined or []],
+    )
+
+    # ping pg (fail fast on a bad connection before starting the server)
+    async with get_async_connection():
+        pass
+
+    try:
+        if transport == "stdio":
+            await mcp_app.run_async(transport=transport)
+        else:
+            await mcp_app.run_async(transport=transport, host=host, port=port)
+    except KeyboardInterrupt:
+        logging.info("KeyboardInterrupt")
+        return 130
+    except Exception:
+        logging.exception("Server terminated with an unexpected error")
+        return 1
+    finally:
+        # Safety net for paths that never reached the server; the app lifespan
+        # already closes the pool on clean shutdown. No-op when no pool is open.
+        try:
+            await close_connection_pool()
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logging.exception("Error during connection pool closure")
+    return None
+
+
+@click.command()
+@click.option(
+    "--overwrite/--no-overwrite",
+    default=True,
+    show_default=True,
+    help="Regenerate the tool files on startup so they reflect the current schema. "
+    "Use --no-overwrite to reuse existing generated files instead.",
+)
+@click.option(
+    "--readonly",
+    is_flag=True,
+    help="Generate only read-only tools (e.g., select functions). No insert or update tools will be generated.",
+)
+@click.option(
+    "--allow-unfiltered",
+    is_flag=True,
+    help="Allow update/delete tools to run without a WHERE filter (affects every row). "
+    "Off by default: unfiltered updates and deletes raise an error.",
+)
+@click.option(
+    "--select-joined",
+    multiple=True,
+    help="Join groups for select, e.g., --select-joined table1,table2,table3",
+)
+@click.option(
+    "--ignore-insert-column",
+    multiple=True,
+    default=["id"],
+    show_default=True,
+    help="Columns to ignore when generating insert functions.",
+)
+@click.option(
+    "--ignore-select-column",
+    multiple=True,
+    default=[],
+    help="Columns to ignore when generating select functions.",
+)
+@click.option(
+    "--ignore-update-column",
+    multiple=True,
+    default=["id"],
+    show_default=True,
+    help="Columns to ignore when generating update functions.",
+)
+@click.option(
+    "--ignore-select-joined-column",
+    multiple=True,
+    default=[],
+    help="Columns to ignore when generating select_joined functions.",
+)
+@click.option(
+    "-t",
+    "--transport",
+    type=click.Choice(["stdio", "http", "sse", "streamable-http"]),
+    default="stdio",
+    show_default=True,
+    help="Transport mechanism for the MCP server.",
+)
+@click.option(
+    "--host",
+    default="127.0.0.1",
+    show_default=True,
+    help="Host to bind the server to (used only with HTTP/SSE transports).",
+)
+@click.option(
+    "--port",
+    default=8000,
+    show_default=True,
+    type=int,
+    help="Port number to bind the server to (used only with HTTP/SSE transports).",
+)
+def main(
+    overwrite: bool,
+    readonly: bool,
+    allow_unfiltered: bool,
+    select_joined: Tuple[str, ...],
+    ignore_insert_column: Tuple[str, ...],
+    ignore_select_column: Tuple[str, ...],
+    ignore_update_column: Tuple[str, ...],
+    ignore_select_joined_column: Tuple[str, ...],
+    transport: TransportName,
+    host: str,
+    port: int,
+) -> None:
+    """Generate scoped PostgreSQL CRUD MCP tools from a live schema."""
+    sys.exit(
+        asyncio.run(
+            runner(
+                overwrite,
+                readonly,
+                allow_unfiltered,
+                select_joined,
+                ignore_insert_column,
+                ignore_select_column,
+                ignore_update_column,
+                ignore_select_joined_column,
+                transport,
+                host,
+                port,
+            )
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

tai42_mcp_dynamic_postgres/config/settings.py

@@ -0,0 +1,40 @@
+from pathlib import Path
+
+from dotenv import load_dotenv
+from pydantic import Field, SecretStr
+from pydantic_settings import BaseSettings
+
+load_dotenv()
+
+
+class PostgresSettings(BaseSettings):
+    model_config = {
+        "env_prefix": "PG_",
+    }
+    tools_dir: str = Field(
+        default_factory=lambda: str(Path.home() / ".cache" / "tai42-mcp-dynamic-postgres" / "tools"),
+        description="Directory for generated tool files",
+        validation_alias="TOOLS_DIR",
+    )
+    host: str = Field("localhost", description="PostgreSQL host")
+    port: int = Field(5432, description="PostgreSQL port")
+    # No defaults: db, user, and password must be supplied; a missing value raises
+    # a ValidationError at startup rather than connecting with a default.
+    db: str = Field(description="Database name")
+    user: str = Field(description="Database user (use a least-privilege role)")
+    password: SecretStr = Field(description="Database password")
+
+    # Statement timeout (milliseconds) applied to every pooled connection. Caps a
+    # runaway query so it cannot hold a pool slot indefinitely. 0 disables it.
+    statement_timeout: int = Field(30000, description="Per-connection statement_timeout in milliseconds (0 disables)")
+
+    # Pool configuration
+    pool_min_size: int = Field(1, description="Minimum number of pooled connections")
+    pool_max_size: int = Field(10, description="Maximum number of pooled connections")
+    pool_timeout: int = Field(10, description="Pool acquire timeout in seconds")
+    pool_max_lifetime: int = Field(300, description="Max lifetime of connection in seconds")
+
+
+# Populated from PG_*/TOOLS_DIR env vars. db/user/password have no defaults, so
+# construction fails loudly if they are missing.
+pg_settings = PostgresSettings()  # pyright: ignore[reportCallIssue]

tai42_mcp_dynamic_postgres/core/app.py

@@ -0,0 +1,23 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
+from fastmcp import FastMCP
+
+from tai42_mcp_dynamic_postgres.database.connection import close_connection_pool, get_connection_pool
+
+
+@asynccontextmanager
+async def lifespan(server: FastMCP) -> AsyncGenerator[None, None]:
+    """Open the connection pool on startup and close it on shutdown.
+
+    Bound to the app lifespan, so the pool closes whenever the server stops,
+    whatever entry point (CLI, embedding host, tests) started it.
+    """
+    await get_connection_pool()
+    try:
+        yield
+    finally:
+        await close_connection_pool()
+
+
+mcp_app: FastMCP = FastMCP(lifespan=lifespan)

tai42_mcp_dynamic_postgres/database/connection.py

@@ -0,0 +1,130 @@
+import asyncio
+import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import Any, Optional
+
+from psycopg import AsyncConnection, AsyncCursor
+from psycopg.conninfo import make_conninfo
+from psycopg.rows import AsyncRowFactory
+from psycopg_pool import AsyncConnectionPool
+
+from tai42_mcp_dynamic_postgres.config.settings import pg_settings
+from tai42_mcp_dynamic_postgres.database.helpers import register_types_loaders
+
+logger = logging.getLogger(__name__)
+
+# The single process-wide pool. Built lazily on first use and closed once on
+# shutdown; never rebuilt merely to be closed.
+_pool: Optional[AsyncConnectionPool[AsyncConnection[Any]]] = None
+_pool_lock = asyncio.Lock()
+
+
+def _build_conninfo() -> str:
+    """Assemble the libpq conninfo string from settings.
+
+    ``make_conninfo`` escapes each value so passwords or names with special
+    characters cannot break out. Built lazily and never stored at module scope, so
+    the cleartext password is not a global.
+    """
+    return make_conninfo(
+        host=pg_settings.host,
+        port=pg_settings.port,
+        dbname=pg_settings.db,
+        user=pg_settings.user,
+        password=pg_settings.password.get_secret_value(),
+        # Cap any single statement so a runaway query cannot hold a pool slot.
+        options=f"-c statement_timeout={pg_settings.statement_timeout}",
+    )
+
+
+async def get_connection_pool() -> AsyncConnectionPool[AsyncConnection[Any]]:
+    global _pool
+    async with _pool_lock:
+        if _pool is None:
+            pool: AsyncConnectionPool[AsyncConnection[Any]] = AsyncConnectionPool(
+                conninfo=_build_conninfo(),
+                min_size=pg_settings.pool_min_size,
+                max_size=pg_settings.pool_max_size,
+                timeout=pg_settings.pool_timeout,
+                max_lifetime=pg_settings.pool_max_lifetime,
+                # Register type loaders once per physical connection, not per checkout.
+                configure=register_types_loaders,
+                open=False,
+            )
+            try:
+                await pool.open()
+            except BaseException:
+                # Opening spun up background workers; close the partial pool so
+                # they are not orphaned, then propagate (``_pool`` stays None).
+                await pool.close()
+                raise
+            _pool = pool
+        return _pool
+
+
+async def close_connection_pool() -> None:
+    # Only close a pool that was actually created; do not build one just to close it.
+    global _pool
+    async with _pool_lock:
+        if _pool is not None:
+            await _pool.close()
+            _pool = None
+
+
+@asynccontextmanager
+async def get_async_connection() -> AsyncGenerator[AsyncConnection[Any], None]:
+    """Async context manager yielding a pooled DB connection in autocommit mode.
+
+    The read/introspection path: autocommit leaves a SELECT-only connection IDLE,
+    so returning it triggers no reset warning or rollback. Autocommit is reset to
+    the pool default before return, or a later write checkout of the same physical
+    connection would inherit it and lose its transactional guarantee.
+    """
+    pool: Optional[AsyncConnectionPool[AsyncConnection[Any]]] = None
+    conn: Optional[AsyncConnection[Any]] = None
+
+    try:
+        pool = await get_connection_pool()
+        conn = await pool.getconn()
+        await conn.set_autocommit(True)
+        yield conn
+    except Exception:
+        logger.exception("Error acquiring or using a pooled database connection")
+        raise
+    finally:
+        if conn is not None and pool is not None:
+            # Reset autocommit to the pool default so a later write checkout stays
+            # transactional. A connection that died mid-statement raises here; log
+            # but fall through to putconn (which discards a broken connection) and
+            # never let this mask the original read error.
+            try:
+                await conn.set_autocommit(False)
+            except Exception:
+                logger.exception("Failed to reset autocommit before returning connection to the pool")
+            await pool.putconn(conn)
+
+
+@asynccontextmanager
+async def cursor(
+    name: str = "",
+    *,
+    binary: bool = False,
+    row_factory: Optional[AsyncRowFactory[Any]] = None,
+    scrollable: Optional[bool] = None,
+    withhold: bool = False,
+) -> AsyncGenerator[AsyncCursor[Any], None]:
+    # Forward row_factory only when supplied so the connection's default
+    # row factory is used otherwise.
+    extra: dict[str, Any] = {} if row_factory is None else {"row_factory": row_factory}
+    async with (
+        get_async_connection() as conn,
+        conn.cursor(
+            name=name,
+            binary=binary,
+            scrollable=scrollable,
+            withhold=withhold,
+            **extra,
+        ) as cur,
+    ):
+        yield cur

tai42_mcp_dynamic_postgres/database/helpers.py

@@ -0,0 +1,44 @@
+import warnings
+from typing import Any, List
+
+from psycopg import AsyncConnection
+from psycopg.abc import Buffer
+from psycopg.adapt import Loader
+from psycopg.types import TypeInfo
+from psycopg.types.json import JsonDumper
+
+
+class VectorLoader(Loader):
+    def load(self, data: Buffer) -> List[float]:
+        s = bytes(data).decode("utf-8")
+        if s.startswith("[") and s.endswith("]"):
+            values = s[1:-1].split(",")
+            return [float(x.strip()) for x in values if x.strip()]
+        else:
+            raise ValueError(f"Invalid vector format: {s}")
+
+
+async def register_vector_as_list(conn: AsyncConnection[Any]) -> None:
+    tinfo = await TypeInfo.fetch(conn, "vector")
+    if not tinfo:
+        warnings.warn(
+            "Postgres type 'vector' not found. Did you enable the pgvector extension?",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+    else:
+        conn.adapters.register_loader(tinfo.oid, VectorLoader)
+
+
+def register_json_dumpers(conn: AsyncConnection[Any]) -> None:
+    # Safety net so a bare ``dict`` binds as JSON. Generated insert/update tools
+    # already wrap json/jsonb values in ``Json(...)`` by column type, so array
+    # columns keep their native list adapter.
+    conn.adapters.register_dumper(dict, JsonDumper)
+
+
+async def register_types_loaders(conn: AsyncConnection[Any]) -> None:
+    # Temporal and uuid types keep psycopg's native loaders, matching the native
+    # annotations the code generator emits for those columns.
+    await register_vector_as_list(conn)
+    register_json_dumpers(conn)

tai42_mcp_dynamic_postgres/docs/index.mdx

@@ -0,0 +1,127 @@
+---
+title: Dynamic Postgres MCP server
+description: Configure and run the PostgreSQL MCP server — the PG_ and TOOLS_DIR settings, the CLI scoping flags, transport selection, and the generated CRUD tools.
+---
+
+Point the server at a PostgreSQL database and it introspects the schema and
+generates one typed MCP tool per DML operation per table. What the agent can reach
+is set two ways that must agree: the CLI flags below decide which tools exist, and
+the PostgreSQL role you connect as decides what those tools can physically do.
+Connect as a dedicated least-privilege role, never a superuser.
+
+## Connection settings
+
+Connection and pooling come from environment variables. `PG_DB`, `PG_USER`, and
+`PG_PASSWORD` have no default — a missing one fails at startup rather than
+connecting with a phantom value.
+
+| Variable | Default | Effect |
+| --- | --- | --- |
+| `PG_HOST` | `localhost` | PostgreSQL host. |
+| `PG_PORT` | `5432` | PostgreSQL port. |
+| `PG_DB` | required | Database name. |
+| `PG_USER` | required | Database user; use a least-privilege role. |
+| `PG_PASSWORD` | required | Database password. Held as a secret. |
+| `PG_STATEMENT_TIMEOUT` | `30000` | Per-connection `statement_timeout` in ms (`0` disables), capping a runaway query. |
+| `PG_POOL_MIN_SIZE` | `1` | Minimum pooled connections. |
+| `PG_POOL_MAX_SIZE` | `10` | Maximum pooled connections. |
+| `PG_POOL_TIMEOUT` | `10` | Pool acquire timeout (seconds). |
+| `PG_POOL_MAX_LIFETIME` | `300` | Maximum connection lifetime (seconds). |
+| `TOOLS_DIR` | `~/.cache/tai42-mcp-dynamic-postgres/tools` | Where generated tool files are written. |
+
+## Scoping flags
+
+The flags decide which tools are generated and how permissive they are.
+
+| Flag | Default | Effect |
+| --- | --- | --- |
+| `--overwrite / --no-overwrite` | on | Regenerate tool files on startup so they track the schema; `--no-overwrite` reuses existing files. |
+| `--readonly` | off | Generate only `select`/`select_joined` tools — no insert, update, or delete. |
+| `--allow-unfiltered` | off | Allow `update`/`delete` to run without a `WHERE` filter (affects every row). Off means unfiltered writes raise. |
+| `--select-joined a,b,c` | — | Generate a joined select over the given tables (repeatable). |
+| `--ignore-insert-column` | `id` | Column to exclude from insert inputs (repeatable). |
+| `--ignore-update-column` | `id` | Column to exclude from update inputs (repeatable). |
+| `--ignore-select-column` | — | Column to exclude from select output models (repeatable). |
+| `--ignore-select-joined-column` | — | Column to exclude from joined select models (repeatable). |
+| `-t, --transport` | `stdio` | `stdio`, `http`, `sse`, or `streamable-http`. |
+| `--host` | `127.0.0.1` | Bind host (HTTP/SSE transports only). |
+| `--port` | `8000` | Bind port (HTTP/SSE transports only). |
+
+## Transport
+
+The default transport is `stdio`: an MCP client launches the server as a
+subprocess and speaks over stdin/stdout. Wire it into a client's server list,
+passing the connection settings through `env`:
+
+```json
+{
+  "mcpServers": {
+    "postgres": {
+      "command": "uvx",
+      "args": ["tai42-mcp-dynamic-postgres", "--readonly"],
+      "env": { "PG_DB": "dbname", "PG_USER": "agent", "PG_PASSWORD": "secret" }
+    }
+  }
+}
+```
+
+For a network-mounted server, choose an HTTP transport and bind a host and port
+(`--host` / `--port` apply only here):
+
+```bash
+uvx tai42-mcp-dynamic-postgres -t http --host 0.0.0.0 --port 8000
+```
+
+## Generated tools
+
+For a table `public.events` you get (unless `--readonly`):
+
+- `select_public_events(where, order_by, limit, offset)`
+- `insert_public_events(params, raise_on_conflict)`
+- `update_public_events(data, where)`
+- `delete_public_events(where)`
+
+Column types map to native Python: temporal columns to `datetime`/`date`/`time`,
+`uuid` to `uuid.UUID`, `numeric`/`decimal` to `Decimal`, `json`/`jsonb` to `Any`,
+and array columns to `list[...]`. `insert` returns the table's real primary key (a
+scalar list for a single-column key, a list of lists for a composite key, or the
+affected row count when the table has none); columns with a database default are
+omittable. `order_by` items take an optional `nulls` (`FIRST`/`LAST`).
+
+### Filtering — `WhereFilter`
+
+`select`, `update`, and `delete` accept a `where` argument. Field names must be
+real columns of the table; unknown fields are rejected, and values are always
+bound parameters, so there is no path for SQL injection through field names.
+
+```jsonc
+// Simple field filters (implicitly ANDed)
+{ "status": { "eq": "open" }, "total": { "gte": 100 } }
+
+// Logical composition
+{ "AND": [ { "status": { "eq": "open" } },
+           { "OR": [ { "total": { "gt": 1000 } }, { "starred": { "eq": true } } ] } ] }
+```
+
+Operators: `eq`, `ne`, `gt`, `gte`, `lt`, `lte`, `like`, `not_like`, `ilike`,
+`not_ilike`, `in`, `not_in`, `between`, `is_null`, and `knn` (pgvector). Logical
+keys: `AND`, `OR`, `NOT`.
+
+### Vector search (pgvector)
+
+When a column is a `vector`, filter or order by similarity. Requires the
+`pgvector` extension enabled in the database.
+
+```jsonc
+{ "embedding": { "knn": { "query": [0.1, 0.2, 0.3],
+                          "distance": "cosine",   // l2 | inner_product | cosine
+                          "threshold": 0.5 } } }
+```
+
+## Docker
+
+```bash
+docker build -t tai42-mcp-dynamic-postgres .
+docker run --rm -e PG_HOST=... -e PG_DB=... -e PG_USER=... -e PG_PASSWORD=... \
+  tai42-mcp-dynamic-postgres tai42-mcp-dynamic-postgres --readonly
+```

tai42_mcp_dynamic_postgres/gen/builders/base_gen.py

@@ -0,0 +1,82 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Tuple
+
+from tai42_mcp_dynamic_postgres.config.settings import pg_settings
+from tai42_mcp_dynamic_postgres.gen.schema.introspect import ColumnInfo, ForeignKey, TableInfo
+
+OUTPUT_DIR = Path(pg_settings.tools_dir)
+OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+TOOLS_SUFFIX = "tools"
+
+# (model_code, tool_code) for one generated tool; model_code may be empty.
+Chunk = Tuple[str, str]
+
+
+class BaseGen(ABC):
+    def __init__(self, prefix: str, imports: str, template: str, ignore_columns: Optional[List[str]] = None) -> None:
+        self.prefix = prefix
+        self.imports = imports
+        self.template = template
+        self.ignore_columns = ignore_columns or []
+
+    @abstractmethod
+    def tool_chunks(self, tables: Dict[str, TableInfo], fks: List[ForeignKey]) -> Iterable[Chunk]:
+        """Yield a ``(model_code, tool_code)`` pair per generated tool."""
+
+    @property
+    def output_path(self) -> Path:
+        return OUTPUT_DIR / f"{self.prefix}_{TOOLS_SUFFIX}.py"
+
+    @property
+    def module_name(self) -> str:
+        return f"{self.prefix}_{TOOLS_SUFFIX}"
+
+    @property
+    def is_exists(self) -> bool:
+        return self.output_path.exists()
+
+    def func_name(self, table: str) -> str:
+        return f"{self.prefix}_{table.replace('.', '_')}"
+
+    def generate_tools(self, tables: Dict[str, TableInfo], fks: List[ForeignKey]) -> List[str]:
+        chunks = [self.imports]
+        for model_code, tool_code in self.tool_chunks(tables, fks):
+            chunks.append(model_code)
+            chunks.append(tool_code)
+        return chunks
+
+    def generate_file(self, tables: Dict[str, TableInfo], fks: List[ForeignKey]) -> None:
+        # Generate fully before truncating the target, so a generation failure
+        # leaves any existing file intact rather than half-written.
+        code = "".join(chunk for chunk in self.generate_tools(tables, fks) if chunk)
+        self.output_path.write_text(code)
+
+
+class TableGen(BaseGen):
+    """Base for generators that emit one tool per table."""
+
+    # Write generators (insert/update/delete) override this: views and
+    # materialized views are read-only and get no write tools.
+    writable_only: bool = False
+
+    def tool_chunks(self, tables: Dict[str, TableInfo], fks: List[ForeignKey]) -> Iterable[Chunk]:
+        for table_info in tables.values():
+            if self.writable_only and not table_info.writable:
+                continue
+            chunk = self.generate_tool(table_info)
+            if chunk is not None:
+                yield chunk
+
+    @abstractmethod
+    def generate_tool(self, table_info: TableInfo) -> Optional[Chunk]:
+        """Build the (model_code, tool_code) for a single table, or None to skip it."""
+
+    @staticmethod
+    def col_names(table_info: TableInfo) -> List[str]:
+        return [col.name for col in table_info.columns]
+
+    def included(self, table_info: TableInfo) -> List[ColumnInfo]:
+        """Columns kept after applying ``ignore_columns``."""
+        return [col for col in table_info.columns if col.name not in self.ignore_columns]