diff --git a/CHANGELOG.md b/CHANGELOG.md index e16de3f..27cc063 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,22 @@ adheres to [Semantic Versioning](https://semver.org/). ## [Unreleased] +### Changed + +- **De-vendored the SQL-safety kernel (roadmap 18.1).** The formerly-vendored + `crystaldba/postgres-mcp` `sql/` subpackage (`src/mcpg/_vendor/`) is + replaced by a first-party `src/mcpg/sql/` package — `allowlist.py` (the + permitted statement/node/function/extension policy, as data), `safety.py` + (`SafeSqlDriver`, the `pglast` allowlist validator), and `driver.py` + (`SqlDriver` / `DbConnPool` / `obfuscate_password`). Behaviour is + **identical** (proven by a differential parity harness — 0 divergence — + plus the ported 760-LOC adversarial suite, a fuzz pass, and a + `/security-review` with no findings); the public seam is unchanged, so no + tool signatures moved. The kernel is now inside the coverage gate, + `mypy --strict`, `ruff`, and `bandit` (the vendored code was excluded from + all four). MCPg now ships no vendored runtime code. Supersedes ADR-0001 + with ADR-0007. + ### Added - **Three more built-in NL→SQL providers (19 → 22).** `translate_nl_to_sql` diff --git a/CLAUDE.md b/CLAUDE.md index a13ed77..4ee5071 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -48,16 +48,27 @@ a read-only HF Spaces demo. `GITHUB_TOKEN`, `DEEPINFRA_TOKEN`). The provider table in `docs/user-guide.md` mirrors it. -## The vendored SQL-safety kernel +## The SQL-safety kernel (first-party) -`src/mcpg/_vendor/sql/` is copied near-verbatim from -[`crystaldba/postgres-mcp`](https://github.com/crystaldba/postgres-mcp) -(MIT, pinned commit `07eb329`; see ADR-0001 and -`src/mcpg/_vendor/README.md`): the `pglast` allowlist validator -(`safe_sql.py`), async pool/driver (`sql_driver.py`), param binding. -It's excluded from the coverage gate and `mypy --strict`; its own -adversarial tests live in `tests/vendor/sql/`. **Never hand-edit** — -local mods are tracked in the vendor README and re-applied on re-sync. +`src/mcpg/sql/` is MCPg's own SQL-safety kernel (roadmap 18.1 de-vendored +the former `crystaldba/postgres-mcp` copy; see ADR-0007, which supersedes +ADR-0001). Three modules, policy separated from mechanism: + +- `sql/allowlist.py` — the SQL-safety **policy as data**: permitted + statement / `pglast` AST node / function / extension sets. The single + auditable decision surface. +- `sql/safety.py` — `SafeSqlDriver`: the `pglast` parse + AST-walker + + read-only execute path. Reads policy from `allowlist.py`; can't widen it. +- `sql/driver.py` — `SqlDriver` / `DbConnPool` / `obfuscate_password` + (pool + execution + credential redaction; no policy). + +Exports `SqlDriver`, `SafeSqlDriver`, `DbConnPool`, `obfuscate_password` +(the seam ~74 `mcpg.*` modules depend on). Fully inside the coverage gate ++ `mypy --strict` + `ruff` + `bandit` like the rest of `mcpg`. Adversarial ++ fuzz tests: `tests/unit/test_sql_kernel_*.py`. Changing the allowlist is +a **security-sensitive** edit — run the adversarial suite and treat the +`docs/reviews/devendor-sql-kernel-security-review.md` threat model as the +bar. ## Workflow conventions diff --git a/NOTICE b/NOTICE index a362d7c..500f08a 100644 --- a/NOTICE +++ b/NOTICE @@ -6,19 +6,20 @@ the repository root for the full text. (Relicensed from AGPL-3.0-or-later in v0.5.x.) ------------------------------------------------------------------------ -Third-party vendored code +Attribution — SQL-safety kernel lineage ------------------------------------------------------------------------ -This product includes software developed by Crystal Corp. +MCPg's SQL-safety kernel (`src/mcpg/sql/`) is first-party code, but it was +re-authored from — and its allowlist policy derives from — the MIT-licensed +`sql/` subpackage of: - Component : PostgreSQL SQL-safety kernel (the `sql/` subpackage) - Location : src/mcpg/_vendor/sql/ - Source : https://github.com/crystaldba/postgres-mcp - Commit : 07eb329c8c48e49640e0d1b5b35465d4d024c3ee - License : MIT License, Copyright (c) 2025, Crystal Corp. - See src/mcpg/_vendor/LICENSE for the full license text. + Project : crystaldba/postgres-mcp ("Postgres MCP Pro") + Source : https://github.com/crystaldba/postgres-mcp + Origin : commit 07eb329c8c48e49640e0d1b5b35465d4d024c3ee (Jan 2026) + License : MIT License, Copyright (c) 2025, Crystal Corp. -The MIT-licensed code above is incorporated under the terms of the MIT -License; both this project and the vendored kernel are now MIT-licensed. -See src/mcpg/_vendor/README.md for provenance details and the re-sync -procedure. +The original kernel was vendored under src/mcpg/_vendor/sql/ from v0.x +until roadmap 18.1 replaced it with the first-party implementation (see +docs/adr/0007-first-party-sql-kernel.md, which supersedes ADR-0001). This +attribution acknowledges that MIT-licensed lineage; MCPg itself remains +MIT-licensed. diff --git a/README.md b/README.md index d34790c..5d4528c 100644 --- a/README.md +++ b/README.md @@ -590,9 +590,9 @@ checklist. ## License -MIT — see [`LICENSE`](https://github.com/devopam/MCPg/blob/main/LICENSE). The vendored SQL-safety kernel at -`src/mcpg/_vendor/sql/` is also MIT-licensed; see [`NOTICE`](https://github.com/devopam/MCPg/blob/main/NOTICE) -for provenance. +MIT — see [`LICENSE`](https://github.com/devopam/MCPg/blob/main/LICENSE). The SQL-safety kernel +(`src/mcpg/sql/`) is first-party, re-authored from the MIT-licensed +`crystaldba/postgres-mcp`; see [`NOTICE`](https://github.com/devopam/MCPg/blob/main/NOTICE) for the lineage. ### Wrapped extensions — licenses you should know about diff --git a/SECURITY.md b/SECURITY.md index c3f1f8e..c859649 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -34,7 +34,8 @@ release notes unless they prefer otherwise. **In scope:** -- The MCPg server code under `src/mcpg/` (excluding `src/mcpg/_vendor/`) +- The MCPg server code under `src/mcpg/` — including the first-party + SQL-safety kernel `src/mcpg/sql/` (the `pglast` allowlist validator) - Authentication & authorisation paths (bearer-token, OIDC, multi-tenancy `SET LOCAL ROLE`) - The capability gates that restrict tool surfaces by access mode @@ -45,9 +46,6 @@ release notes unless they prefer otherwise. **Out of scope:** - Vulnerabilities in PostgreSQL itself -- Vulnerabilities in the vendored SQL-safety kernel at - `src/mcpg/_vendor/sql/` — those go upstream to - `crystaldba/postgres-mcp` - Issues that require an attacker to already have `unrestricted` access mode AND `MCPG_ALLOW_DDL=true` (that combination is by-design root access) diff --git a/docs/adr/0001-build-approach.md b/docs/adr/0001-build-approach.md index 1ae7871..cfa4a91 100644 --- a/docs/adr/0001-build-approach.md +++ b/docs/adr/0001-build-approach.md @@ -1,8 +1,13 @@ # ADR-0001: Build approach — adopt `crystaldba/postgres-mcp` as a hard-forked base -- **Status:** accepted +- **Status:** superseded by [ADR-0007](0007-first-party-sql-kernel.md) (2026-07-08) - **Date:** 2026-05-20 +> **Superseded.** The vendored SQL-safety kernel this ADR introduced was +> replaced by a first-party implementation in roadmap 18.1 — see +> [ADR-0007](0007-first-party-sql-kernel.md). The historical rationale +> below is retained as the record of the original decision. + ## Context The official `@modelcontextprotocol/server-postgres` is deprecated and archived diff --git a/docs/adr/0007-first-party-sql-kernel.md b/docs/adr/0007-first-party-sql-kernel.md new file mode 100644 index 0000000..c63ca06 --- /dev/null +++ b/docs/adr/0007-first-party-sql-kernel.md @@ -0,0 +1,81 @@ +# ADR-0007: First-party SQL-safety kernel (de-vendor `crystaldba/postgres-mcp`) + +- **Status:** accepted +- **Date:** 2026-07-08 +- **Supersedes:** [ADR-0001](0001-build-approach.md) (the hard-fork / vendor decision) + +## Context + +[ADR-0001](0001-build-approach.md) chose to hard-fork MCPg from +`crystaldba/postgres-mcp` and **vendor** its `sql/` subpackage — the +`pglast`-AST SQL-safety allowlist validator plus the async connection +pool / driver — near-verbatim under `src/mcpg/_vendor/sql/` (MIT, pinned +commit `07eb329`). That was the right call to ship fast, but it left two +standing costs: + +1. **~1.3k lines of security-critical code outside every quality gate.** + The vendored kernel — including the SQL allowlist itself — was excluded + from the coverage gate, `mypy --strict`, `ruff`, and `bandit` + (`pyproject.toml` carve-outs). +2. **Manual re-syncs.** Upstream changes had to be re-applied by hand with + local modifications tracked in a vendor README. + +It was also the last third-party runtime *code* MCPg shipped. Roadmap 18.1 +set out to own it. + +## Decision + +Replace the vendored kernel with a **first-party** `src/mcpg/sql/` package, +re-authored from the MIT-licensed original (attribution retained in +`NOTICE`). Recon showed three of the vendored modules (`bind_params.py`, +`extension_utils.py`, `index.py`) had no consumer outside `_vendor/`, so +they were **deleted, not rebuilt**. The kernel is re-architected to +**separate policy from mechanism**: + +- `sql/allowlist.py` — the permitted statement / AST-node / function / + extension sets, as **data** (the single auditable decision surface). +- `sql/safety.py` — `SafeSqlDriver`: the `pglast` parse + AST-walker + + read-only execute path. Reads policy from `allowlist.py`; cannot widen it. +- `sql/driver.py` — `SqlDriver` / `DbConnPool` / `obfuscate_password` + (pool + execution + credential redaction; no policy). + +The public seam (`SqlDriver`, `SafeSqlDriver`, `DbConnPool`, +`obfuscate_password`) is unchanged, so the ~74 consuming modules only +changed an import path. + +**Faithful re-author, not a redesign.** The security *behaviour* is pinned +identical to the vendored validator: + +- A **differential parity harness** ran a safe/unsafe/malformed corpus + through both validators — **0 divergence**. +- The 760-LOC adversarial suite (ported) passes 100% against `mcpg.sql`. +- A **fuzz pass** confirms the validator only ever returns a clean verdict + (accept / `ValueError`) — never a crash / hang — on adversarial input. +- A dedicated security-review gate (allowlist audit, `/security-review`, + threat model) passed — see + [`../reviews/devendor-sql-kernel-security-review.md`](../reviews/devendor-sql-kernel-security-review.md). + +## Consequences + +- The SQL-safety kernel is now inside the coverage gate (90%), + `mypy --strict`, `ruff`, and `bandit` — a net security improvement over + the vendored state. +- `src/mcpg/_vendor/` and `tests/vendor/` are deleted; the five + `pyproject.toml` `_vendor` carve-outs are removed; the module map and + docs (`architecture.md`, `CLAUDE.md`, `SECURITY.md`, `NOTICE`) are + first-party. +- MCPg ships **no vendored runtime code**. (`pglast`, `psycopg`, and + `psycopg-pool` remain upstream *library* dependencies — not vendored + source.) +- Attribution to `crystaldba/postgres-mcp` (MIT) is retained in `NOTICE` + as the design lineage; the in-query `/* crystaldba */` marker is kept + for now (a cosmetic rename is a possible follow-up). + +## Alternatives considered + +- **Keep vendoring.** Rejected — perpetuates the gate carve-outs and the + manual re-sync burden; ADR-0001's "ship fast" rationale no longer applies. +- **Redesign the validator (stricter/simpler allowlist).** Deferred — a + behaviour change would need its own threat model + review. This ADR is a + faithful re-author; any policy tightening is a separate, reviewed change + (see the pre-existing observations in the security-review note). diff --git a/docs/architecture.md b/docs/architecture.md index 3a88970..ee93505 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -51,7 +51,7 @@ flowchart TD runs the SQL and maps rows to typed dataclasses. 4. The driver stack decides exactly which pool the SQL hits: - **SafeSqlDriver** — agent-supplied SQL is parsed and - allowlisted via the vendored kernel before execution. + allowlisted via `mcpg.sql` (the first-party kernel) before execution. - **RoutedSqlDriver** — when `MCPG_REPLICA_URLS` is set, `force_readonly=True` queries round-robin across healthy replicas; writes always go to the primary. @@ -64,7 +64,7 @@ flowchart TD --- -## Module map (104 modules) +## Module map (103 modules) Every `mcpg.*` module and what it owns, alphabetical. The layered request path through these lives in the [Overview](#overview) diagram; @@ -73,7 +73,6 @@ this table is the exhaustive index. Regenerate with | Module | Responsibility | |---|---| -| `mcpg._vendor` | Vendored MIT-licensed `SafeSqlDriver` + connection-pool kernel (SQL parse / allowlist / bind). | | `mcpg.about` | MCPg self-description. | | `mcpg.advisors` | Schema advisors — codified lint rules over the PG catalog. | | `mcpg.aio` | `AIO` — PG 19 asynchronous-I/O subsystem coverage. | @@ -184,17 +183,28 @@ this table is the exhaustive index. Regenerate with --- -## The vendored SQL-safety kernel - -`src/mcpg/_vendor/sql/` is a pinned copy of the SQL-safety -subpackage from -[`crystaldba/postgres-mcp`](https://github.com/crystaldba/postgres-mcp) -(MIT). It provides `SafeSqlDriver` — a `pglast`-AST allowlist -validator — and the base connection pool / driver. - -The kernel is kept near-verbatim, excluded from the coverage gate -and `mypy`, and re-synced via the procedure in -`src/mcpg/_vendor/README.md`. See [ADR-0001](adr/0001-build-approach.md). +## The SQL-safety kernel (first-party) + +`src/mcpg/sql/` is MCPg's own SQL-safety kernel, split into policy +and mechanism: + +- `sql/allowlist.py` — the permitted statement / `pglast` AST node / + function / extension sets, **as data** (the single auditable policy + surface). +- `sql/safety.py` — `SafeSqlDriver`: the `pglast` parse + AST-walker + + read-only execute path. Reads policy from `allowlist.py`; the + walker can't widen it. +- `sql/driver.py` — `SqlDriver` / `DbConnPool` / `obfuscate_password` + (pool + execution + credential redaction; no policy). + +It's fully inside the coverage gate + `mypy --strict` + `ruff` + +`bandit`. Adversarial + fuzz tests live in +`tests/unit/test_sql_kernel_*.py`; the threat model and security +sign-off are in +[`reviews/devendor-sql-kernel-security-review.md`](reviews/devendor-sql-kernel-security-review.md). +It was de-vendored from `crystaldba/postgres-mcp` (MIT) — see +[ADR-0007](adr/0007-first-party-sql-kernel.md), which supersedes +[ADR-0001](adr/0001-build-approach.md). --- @@ -294,14 +304,13 @@ otherwise. ## Testing approach -MCPg is test-driven across four suites: +MCPg is test-driven across three suites: | Suite | Scope | |---|---| -| `tests/unit/` | Fake-driver tests with a 90% coverage gate. Authored code only — the vendored kernel keeps its own tests. | +| `tests/unit/` | Fake-driver tests with a 90% coverage gate — all first-party code, including the SQL-safety kernel (`test_sql_kernel_*.py`: the adversarial allowlist suite + a fuzz/robustness pass). | | `tests/integration/` | Real PostgreSQL — requires `MCPG_TEST_DATABASE_URL`. CI runs the matrix against PostgreSQL 14, 15, 16, 17, 18 on every push (pgvector + PostGIS + AGE image), plus an **experimental PG 19** lane (pgvector built from source against `postgres:19beta1`) and a **WarehousePG** (Greenplum-derived MPP) characterisation lane. | | `tests/contract/` | Tool-surface snapshot (`tool_surface.snapshot.json`) + the doc-table drift-guard (`test_doc_tables.py`), so a new tool or module can't ship undocumented. | -| `tests/vendor/` | The vendored kernel's own upstream tests, kept for adversarial SQL-injection coverage. | The integration container is built from `.github/ci-postgres.Dockerfile` and includes `pgvector`, diff --git a/docs/contributing/adding-tools.md b/docs/contributing/adding-tools.md index ff86330..0f63fe9 100644 --- a/docs/contributing/adding-tools.md +++ b/docs/contributing/adding-tools.md @@ -41,7 +41,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver from mcpg.extensions import extension_installed # Module-level constants (mode allowlists, defaults). Inline so the diff --git a/docs/feature-shortlist.md b/docs/feature-shortlist.md index a5e06ee..f95bfc4 100644 --- a/docs/feature-shortlist.md +++ b/docs/feature-shortlist.md @@ -270,19 +270,15 @@ empty scratch database), which shows off none of the tool surface. ## 18. De-vendor the SQL-safety kernel — build our own -MCPg's SQL-safety core is currently **vendored** from a third party -(see [ADR-0001](adr/) and [`src/mcpg/_vendor/README.md`](../src/mcpg/_vendor/README.md)): -the `pglast`-based allowlist validator, async connection pool / driver, -and parameter binding all live under `src/mcpg/_vendor/sql/`, copied -near-verbatim from -[`crystaldba/postgres-mcp`](https://github.com/crystaldba/postgres-mcp) -(MIT, pinned commit `07eb329`, Copyright © 2025 Crystal Corp). The goal -is to **own this layer outright** — remove the upstream dependency and -replace it with a first-party implementation MCPg fully controls. +MCPg's SQL-safety core was **vendored** from `crystaldba/postgres-mcp` +(MIT) under `src/mcpg/_vendor/sql/`. Roadmap 18.1 replaced it with a +first-party implementation MCPg fully controls — removing the last +vendored runtime code and folding the SQL allowlist into every quality +gate. | # | Item | Effort | Value | Notes | |---|---|---|---|---| -| 18.1 | ⬜ **Planned (agenda, no commitment yet).** **Replace the vendored `crystaldba/postgres-mcp` SQL-safety kernel with a first-party implementation.** In scope: `safe_sql.py` (the `pglast` parse + statement/keyword allowlist that every agent-supplied query passes through), `sql_driver.py` (the async psycopg pool + `DbConnPool`, incl. MCPg's ADR-0003 pool-sizing local mod), `bind_params.py`, `extension_utils.py`, `index.py`. Constraints: (a) **no safety regression** — the adversarial SQL-injection suite in `tests/vendor/sql/test_safe_sql.py` must be preserved (or re-implemented) and kept green, and the tool-surface contract test must not change; (b) preserve the `SafeSqlDriver` / pool public API the rest of `mcpg.*` depends on, or migrate call sites in the same effort; (c) fold the code into the coverage gate + `mypy --strict` once it's ours (the `_vendor/` carve-outs go away); (d) drop the vendored dir, its `README.md`/`LICENSE`, and the ADR-0001 rationale, and record the change in a new ADR. Likely a multi-PR effort: first re-home + test-parity, then simplify/own. **Phased plan: [`plans/devendor-sql-kernel.md`](plans/devendor-sql-kernel.md)** — recon found `bind_params.py` / `extension_utils.py` / `index.py` are **dead** (no consumer outside `_vendor/`), so they're deleted not rebuilt; the real rebuild is just `safe_sql.py` + `sql_driver.py`, and the consumer seam is 4 names (`SqlDriver` / `SafeSqlDriver` / `DbConnPool` / `obfuscate_password`). | L | Medium | Removes the last third-party runtime core; full control over the safety semantics. Revisit ADR-0001's "why vendor" reasoning when we start. | +| 18.1 | ✅ **Shipped.** Replaced the vendored `crystaldba/postgres-mcp` kernel with a first-party `src/mcpg/sql/` package — `allowlist.py` (policy as data), `safety.py` (`SafeSqlDriver` walker), `driver.py` (`SqlDriver`/`DbConnPool`/`obfuscate_password`). The dead `bind_params.py` / `extension_utils.py` / `index.py` were deleted, not rebuilt. All ~74 consumers swung to `mcpg.sql`; `_vendor/` + `tests/vendor/` deleted; the 5 quality-gate carve-outs removed (the kernel is now inside coverage + `mypy --strict` + `ruff` + `bandit`). Faithful re-author, proven identical by a differential parity harness (0 divergence) + the ported 760-LOC adversarial suite + a fuzz pass + `/security-review` (no findings). Supersedes ADR-0001 with [ADR-0007](adr/0007-first-party-sql-kernel.md); plan in [`plans/devendor-sql-kernel.md`](plans/devendor-sql-kernel.md), sign-off in [`reviews/devendor-sql-kernel-security-review.md`](reviews/devendor-sql-kernel-security-review.md). | L | Medium | Removes the last third-party runtime core; full control over the safety semantics. | --- diff --git a/docs/release-process.md b/docs/release-process.md index 89b75d8..3f508f7 100644 --- a/docs/release-process.md +++ b/docs/release-process.md @@ -144,7 +144,7 @@ description = "A production-grade PostgreSQL Model Context Protocol (MCP) server readme = "README.md" requires-python = ">=3.12" license = "MIT" -license-files = ["LICENSE", "src/mcpg/_vendor/LICENSE"] +license-files = ["LICENSE"] ``` ### 3.2 What to add before the first release @@ -160,7 +160,7 @@ description = "A production-grade PostgreSQL Model Context Protocol (MCP) server readme = "README.md" # → "Project description" tab requires-python = ">=3.12" license = "MIT" -license-files = ["LICENSE", "src/mcpg/_vendor/LICENSE"] +license-files = ["LICENSE"] # Surfaces under "Author" on the project page + in `pip show`. authors = [ diff --git a/pyproject.toml b/pyproject.toml index 237deb8..e02f384 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -170,18 +170,16 @@ filterwarnings = [ [tool.coverage.run] source = ["src/mcpg"] -omit = ["src/mcpg/_vendor/*"] [tool.coverage.report] -# Coverage gate applies to authored code only; vendored code is excluded above. +# Coverage gate applies to all authored code — the SQL kernel is first-party. fail_under = 90 show_missing = true [tool.ruff] line-length = 120 target-version = "py312" -# Vendored code is kept verbatim and is not linted/formatted to our standard. -extend-exclude = ["src/mcpg/_vendor", "tests/vendor", "scratch", "demo.py"] +extend-exclude = ["scratch", "demo.py"] # Honour the exclusions even when files are passed explicitly (e.g. pre-commit). force-exclude = true @@ -194,12 +192,7 @@ known-first-party = ["mcpg"] [tool.mypy] python_version = "3.14" strict = true -# Vendored code keeps its own (looser) typing; do not gate on it. -exclude = ["src/mcpg/_vendor/", "tests/vendor/", "scratch/"] - -[[tool.mypy.overrides]] -module = ["mcpg._vendor.*"] -ignore_errors = true +exclude = ["scratch/"] [[tool.mypy.overrides]] # pglast ships no type stubs. @@ -225,6 +218,6 @@ module = [ ignore_missing_imports = true [tool.bandit] -exclude_dirs = ["tests", "src/mcpg/_vendor"] +exclude_dirs = ["tests"] skips = ["B101", "B608", "B110"] diff --git a/src/mcpg/_vendor/LICENSE b/src/mcpg/_vendor/LICENSE deleted file mode 100644 index 49eef58..0000000 --- a/src/mcpg/_vendor/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2025, Crystal Corp. - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/src/mcpg/_vendor/README.md b/src/mcpg/_vendor/README.md deleted file mode 100644 index c61c32a..0000000 --- a/src/mcpg/_vendor/README.md +++ /dev/null @@ -1,37 +0,0 @@ -# Vendored code - -This directory contains third-party code vendored into MCPg. It is **not** -authored by the MCPg project and is excluded from the coverage gate and from -`mypy --strict`. - -## `sql/` — PostgreSQL SQL-safety kernel - -- **Source:** [`crystaldba/postgres-mcp`](https://github.com/crystaldba/postgres-mcp) -- **Pinned commit:** `07eb329c8c48e49640e0d1b5b35465d4d024c3ee` (2026-01-22) -- **Licence:** MIT — see [`LICENSE`](LICENSE) (Copyright (c) 2025, Crystal Corp.) -- **What it is:** the `pglast`-based SQL allowlist validator (`safe_sql.py`), - async connection pool / driver (`sql_driver.py`), parameter binding - (`bind_params.py`), and extension utilities. See ADR-0001 for why only this - subpackage was vendored. - -### Local modifications - -The source files are near-verbatim copies. Deliberate local changes: - -- **Test import paths** (`tests/vendor/`): `postgres_mcp.sql` → - `mcpg._vendor.sql`. The `sql/` source files use relative imports and were - otherwise copied unchanged. -- **`sql_driver.py` — `DbConnPool` pool sizing (ADR-0003):** `__init__` gained - `min_size`/`max_size` parameters (defaulting to `1`/`5`, reproducing the - original behaviour), used in `pool_connect`. A ~3-line, behaviour-preserving - change; re-apply on re-sync. Marked in-file with `MCPg local modification`. - -### Re-sync procedure - -To pull upstream security fixes: - -1. `git clone https://github.com/crystaldba/postgres-mcp /tmp/pg-mcp` -2. Diff `/tmp/pg-mcp/src/postgres_mcp/sql/` against `./sql/`. -3. Apply relevant changes; do not introduce new imports outside the subpackage. -4. Update the pinned commit above and note the change in `CHANGELOG.md`. -5. Run `tests/vendor/` to confirm behaviour is preserved. diff --git a/src/mcpg/_vendor/__init__.py b/src/mcpg/_vendor/__init__.py deleted file mode 100644 index 0a67498..0000000 --- a/src/mcpg/_vendor/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Vendored third-party code. See README.md for provenance and licence. - -Code under this package is NOT authored by the MCPg project. It is a pinned -copy of an upstream project, kept verbatim except for import paths. Do not -hand-edit; re-sync via the procedure in README.md. -""" diff --git a/src/mcpg/_vendor/sql/__init__.py b/src/mcpg/_vendor/sql/__init__.py deleted file mode 100644 index 1fded3b..0000000 --- a/src/mcpg/_vendor/sql/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -"""SQL utilities.""" - -from .bind_params import ColumnCollector -from .bind_params import SqlBindParams -from .bind_params import TableAliasVisitor -from .extension_utils import check_extension -from .extension_utils import check_hypopg_installation_status -from .extension_utils import check_postgres_version_requirement -from .extension_utils import get_postgres_version -from .extension_utils import reset_postgres_version_cache -from .index import IndexDefinition -from .safe_sql import SafeSqlDriver -from .sql_driver import DbConnPool -from .sql_driver import SqlDriver -from .sql_driver import obfuscate_password - -__all__ = [ - "ColumnCollector", - "DbConnPool", - "IndexDefinition", - "SafeSqlDriver", - "SqlBindParams", - "SqlDriver", - "TableAliasVisitor", - "check_extension", - "check_hypopg_installation_status", - "check_postgres_version_requirement", - "get_postgres_version", - "obfuscate_password", - "reset_postgres_version_cache", -] diff --git a/src/mcpg/_vendor/sql/bind_params.py b/src/mcpg/_vendor/sql/bind_params.py deleted file mode 100644 index ec6516e..0000000 --- a/src/mcpg/_vendor/sql/bind_params.py +++ /dev/null @@ -1,816 +0,0 @@ -# --- Parameter replacement --- - -import logging -import re -from typing import Any - -from pglast import parse_sql -from pglast.ast import A_Expr -from pglast.ast import ColumnRef -from pglast.ast import JoinExpr -from pglast.ast import Node -from pglast.ast import RangeVar -from pglast.ast import SelectStmt -from pglast.ast import SortBy -from pglast.ast import SortGroupClause -from pglast.visitors import Visitor - -from .safe_sql import SafeSqlDriver -from .sql_driver import SqlDriver - -logger = logging.getLogger(__name__) - - -# --- Visitor Classes --- - - -class TableAliasVisitor(Visitor): - """Extracts table aliases and names from the SQL AST.""" - - def __init__(self) -> None: - super().__init__() # Initialize the base Visitor class - self.aliases: dict[str, str] = {} - self.tables: set[str] = set() - - def __call__(self, node): - super().__call__(node) - return self.aliases, self.tables - - def visit_RangeVar(self, ancestors: list[Node], node: Node) -> None: # noqa: N802 - """Visit table references, including those in FROM clause.""" - if isinstance(node, RangeVar): # Type narrowing for RangeVar - if node.relname is not None: - self.tables.add(node.relname) - if node.alias and node.alias.aliasname is not None: - self.aliases[node.alias.aliasname] = str(node.relname) - - def visit_JoinExpr(self, ancestors: list[Node], node: Node) -> None: # noqa: N802 - """Visit both sides of JOIN expressions.""" - if isinstance(node, JoinExpr): # Type narrowing for JoinExpr - if node.larg is not None: - self(node.larg) # type: ignore - if node.rarg is not None: - self(node.rarg) # type: ignore - - -class ColumnCollector(Visitor): - """ - Collects columns used in WHERE, JOIN, ORDER BY, GROUP BY, HAVING, and SELECT clauses. - With improved handling of column aliases. - """ - - def __init__(self) -> None: - super().__init__() - self.context_stack = [] # Stack of (tables, aliases) for each scope - self.columns = {} # Collected columns, keyed by table - self.target_list = None - self.inside_select = False - self.column_aliases = {} # Track column aliases and their definitions - self.current_query_level = 0 # Track nesting level for subqueries - - def __call__(self, node): - super().__call__(node) - return self.columns - - def visit_SelectStmt(self, ancestors: list[Node], node: Node) -> None: # noqa: N802 - """Visit a SelectStmt node and process its targetList for column aliases.""" - if isinstance(node, SelectStmt): - self.inside_select = True - self.current_query_level += 1 - query_level = self.current_query_level - - # Collect tables and aliases - alias_visitor = TableAliasVisitor() - if hasattr(node, "fromClause") and node.fromClause: - for from_item in node.fromClause: - alias_visitor(from_item) - scope_tables = alias_visitor.tables - scope_aliases = alias_visitor.aliases - - # Push new context for this scope - self.context_stack.append((scope_tables, scope_aliases)) - - # First pass: collect column aliases from targetList - if hasattr(node, "targetList") and node.targetList: - self.target_list = node.targetList - for target_entry in self.target_list: - if hasattr(target_entry, "name") and target_entry.name: - # This is a column alias - col_alias = target_entry.name - # Store the expression node for this alias - if hasattr(target_entry, "val"): - self.column_aliases[col_alias] = { - "node": target_entry.val, - "level": query_level, - } - - # Second pass: process the rest of the query - self._process_query_clauses(node) - - # Pop context after processing - self.context_stack.pop() - self.inside_select = False - self.current_query_level -= 1 - - def _process_query_clauses(self, node): - """Process various query clauses for column collection.""" - # Process targetList expressions - if hasattr(node, "targetList") and node.targetList: - self.target_list = node.targetList - for target_entry in self.target_list: - if hasattr(target_entry, "val"): - self(target_entry.val) - - # Handle GROUP BY clause - if hasattr(node, "groupClause") and node.groupClause: - for group_item in node.groupClause: - if isinstance(group_item, SortGroupClause) and isinstance(group_item.tleSortGroupRef, int): - ref_index = group_item.tleSortGroupRef - if self.target_list and ref_index <= len(self.target_list): - target_entry = self.target_list[ref_index - 1] # 1-based index - if hasattr(target_entry, "val"): - self(target_entry.val) - if hasattr(target_entry, "expr"): - self(target_entry.expr) # Visit the expression - - # Process WHERE clause (including subqueries) - if hasattr(node, "whereClause") and node.whereClause: - self(node.whereClause) - - # Process FROM clause (may contain subqueries) - if hasattr(node, "fromClause") and node.fromClause: - for from_item in node.fromClause: - self(from_item) - - # Process HAVING clause - if hasattr(node, "havingClause") and node.havingClause: - self(node.havingClause) - - # Handle ORDER BY clause - if hasattr(node, "sortClause") and node.sortClause: - for sort_item in node.sortClause: - self._process_sort_item(sort_item) - - def _process_sort_item(self, sort_item): - """Process a sort item, resolving column aliases if needed.""" - if not hasattr(sort_item, "node"): - return - - # If it's a simple column reference, it might be an alias - if isinstance(sort_item.node, ColumnRef) and hasattr(sort_item.node, "fields") and sort_item.node.fields: - fields = [f.sval for f in sort_item.node.fields if hasattr(f, "sval")] - if len(fields) == 1: - col_name = fields[0] - # Check if this is a known alias - if col_name in self.column_aliases: - # Process the original expression instead - alias_info = self.column_aliases[col_name] - if alias_info["level"] == self.current_query_level: - self(alias_info["node"]) - return - - # Regular processing for non-alias sort items - self(sort_item.node) - - def visit_ColumnRef(self, ancestors: list[Node], node: Node) -> None: # noqa: N802 - """Visit a ColumnRef node and collect column names, skipping aliases.""" - if isinstance(node, ColumnRef) and self.inside_select: - if not hasattr(node, "fields") or not node.fields: - return - - fields = [f.sval if hasattr(f, "sval") else "*" for f in node.fields] - - # Skip collecting if this is a reference to a column alias - if len(fields) == 1 and (fields[0] == "*" or fields[0] in self.column_aliases): - return - - if len(fields) == 2 and fields[1] == "*": - return - - # Use current scope's tables and aliases - current_tables, current_aliases = self.context_stack[-1] if self.context_stack else ({}, {}) - - if len(fields) == 2: # Qualified column (e.g., u.name) - table_or_alias, column = fields - table = current_aliases.get(table_or_alias, table_or_alias) - if table not in self.columns: - self.columns[table] = set() - self.columns[table].add(column) - elif len(fields) == 1: # Unqualified column - column = fields[0] - if len(current_tables) == 1: # Only one table in scope - table = next(iter(current_tables)) - if table not in self.columns: - self.columns[table] = set() - self.columns[table].add(column) - else: - # Try to find which table this column belongs to - for table in current_tables: - if self._column_exists(table, column): - if table not in self.columns: - self.columns[table] = set() - self.columns[table].add(column) - break - - def _column_exists(self, table: str, column: str) -> bool: - """Check if column exists in table.""" - # This is a placeholder. We'd query the schema - # Ideally this would be cached to avoid repeated queries - return True # Default to True for the sake of collecting all possibilities - - def visit_A_Expr(self, ancestors: list[Node], node: Node) -> None: # noqa: N802 - """ - Visit an A_Expr node (arithmetic or comparison expression). - """ - if isinstance(node, A_Expr) and self.inside_select: - # Process left expression - if hasattr(node, "lexpr") and node.lexpr: - self(node.lexpr) - if isinstance(node.lexpr, SelectStmt): - alias_visitor = TableAliasVisitor() - alias_visitor(node.lexpr) - self.context_stack.append((alias_visitor.tables, alias_visitor.aliases)) - self(node.lexpr) - self.context_stack.pop() - - # Process right expression - if hasattr(node, "rexpr") and node.rexpr: - if isinstance(node.rexpr, SelectStmt): - alias_visitor = TableAliasVisitor() - alias_visitor(node.rexpr) - self.context_stack.append((alias_visitor.tables, alias_visitor.aliases)) - self(node.rexpr) - self.context_stack.pop() - else: - self(node.rexpr) - - # Special handling for IN clauses with subqueries - if hasattr(node, "kind") and node.kind == 0: # 0 is the kind for IN operator - if hasattr(node, "rexpr") and node.rexpr and isinstance(node.rexpr, SelectStmt): - # Process the subquery in the IN clause - alias_visitor = TableAliasVisitor() - alias_visitor(node.rexpr) - self.context_stack.append((alias_visitor.tables, alias_visitor.aliases)) - self(node.rexpr) - self.context_stack.pop() - - def visit_JoinExpr(self, ancestors: list[Node], node: Node) -> None: # noqa: N802 - """ - Visit a JoinExpr node to handle JOIN conditions. - """ - if isinstance(node, JoinExpr) and self.inside_select: # Type narrowing for JoinExpr - if hasattr(node, "larg") and node.larg: - self(node.larg) - if hasattr(node, "rarg") and node.rarg: - self(node.rarg) - if hasattr(node, "quals") and node.quals: - self(node.quals) - - def visit_SortBy(self, ancestors: list[Node], node: Node) -> None: # noqa: N802 - """ - Visit a SortBy node (ORDER BY expression). - """ - if isinstance(node, SortBy) and self.inside_select: # Type narrowing for SortBy - if hasattr(node, "node") and node.node: - self(node.node) - - -class SqlBindParams: - """ - Replaces parameter placeholders with appropriate values based on column statistics. - """ - - def __init__(self, sql_driver: SqlDriver): - self.sql_driver = sql_driver - self._column_stats_cache = {} - - async def replace_parameters(self, query: str) -> str: - """Replace parameter placeholders with appropriate values based on column statistics. - - This handles queries from pg_stat_statements where literals - have been replaced with $1, $2, etc. - """ - try: - modified_query = query - # Find all parameter placeholders - param_matches = list(re.finditer(r"\$\d+", query)) - if not param_matches: - logger.debug(f"No parameters found for query: {query[:50]}...") - return query - - # Handle common special cases in a specific order to prevent incorrect replacements - - # 1. Handle LIMIT clauses - these should always be replaced with integers - limit_pattern = re.compile(r"limit\s+\$(\d+)", re.IGNORECASE) - modified_query = limit_pattern.sub(r"limit 100", modified_query) - - # 2. Handle static INTERVAL expressions - interval_pattern = re.compile(r"interval\s+'(\d+)\s+([a-z]+)'", re.IGNORECASE) - modified_query = interval_pattern.sub(lambda m: f"interval '2 {m.group(2)}'", modified_query) - - # 3. Handle parameterized INTERVAL expressions (INTERVAL $1) - param_interval_pattern = re.compile(r"interval\s+\$(\d+)", re.IGNORECASE) - modified_query = param_interval_pattern.sub("interval '2 days'", modified_query) - - # 4. Handle OFFSET clauses - similar to LIMIT - offset_pattern = re.compile(r"offset\s+\$(\d+)", re.IGNORECASE) - modified_query = offset_pattern.sub(r"offset 0", modified_query) - - # Find all remaining parameter placeholders - param_matches = list(re.finditer(r"\$\d+", modified_query)) - if not param_matches: - return modified_query - - # Then, handle BETWEEN clauses as special cases - between_pattern = re.compile(r"(\w+(?:\.\w+)?)\s+between\s+\$(\d+)\s+and\s+\$(\d+)", re.IGNORECASE) - for match in between_pattern.finditer(query): - column_ref, param1, param2 = match.groups() - # Extract table and column name from the reference - if "." in column_ref: - parts = column_ref.split(".") - alias = parts[0] - col_name = parts[1] - # Resolve the table alias - table_columns = self.extract_columns(query) - table_name = None - for tbl, _cols in table_columns.items(): - if any(alias == a for a in self._get_table_aliases(query, tbl)): - table_name = tbl - break - else: - # No alias, try to find the column in any table - col_name = column_ref - table_columns = self.extract_columns(query) - table_name = None - for tbl, cols in table_columns.items(): - if col_name in cols: - table_name = tbl - break - - # Default numeric bounds if statistics not available - lower_bound = 10 - upper_bound = 100 - if table_name and col_name: - stats = await self._get_column_statistics(table_name, col_name) - if stats: - # Get appropriate values for both bounds - lower_bound = self._get_bound_values(stats, is_lower=True) - upper_bound = self._get_bound_values(stats, is_lower=False) - - # Replace both parameters in the BETWEEN clause - param1_pattern = r"\$" + param1 - param2_pattern = r"\$" + param2 - modified_query = re.sub(param1_pattern, str(lower_bound), modified_query) - modified_query = re.sub(param2_pattern, str(upper_bound), modified_query) - - # Now handle remaining parameters normally - # Recompute matches after BETWEEN replacements - param_matches = list(re.finditer(r"\$\d+", modified_query)) - if not param_matches: - return modified_query - - table_columns = self.extract_columns(query) - if not table_columns: - return self._replace_parameters_generic(modified_query) - - # Process each remaining parameter - for match in reversed(param_matches): - param_position = match.start() - - # Extract a narrower context - clause_start = max( - modified_query.rfind(" where ", 0, param_position), - modified_query.rfind(" and ", 0, param_position), - modified_query.rfind(" or ", 0, param_position), - modified_query.rfind(",", 0, param_position), - modified_query.rfind("(", 0, param_position), - -1, - ) - - if clause_start == -1: - clause_start = max(0, param_position - 100) - - preceding_text = modified_query[clause_start : param_position + 2] - - # Try to identify which column this parameter belongs to - column_info = self._identify_parameter_column(preceding_text, table_columns) - if column_info: - table_name, column_name = column_info - stats = await self._get_column_statistics(table_name, column_name) - if stats: - replacement = self._get_replacement_value(stats, preceding_text) - else: - replacement = self._get_generic_replacement(preceding_text) - else: - replacement = self._get_generic_replacement(preceding_text) - - modified_query = modified_query[: match.start()] + replacement + modified_query[match.end() :] - - return modified_query - except Exception as e: - raise ValueError("Error replacing parameters") from e - - def _get_bound_values(self, stats: dict[str, Any], is_lower: bool = True) -> Any: - """Get appropriate bound values for range queries based on column statistics. - - Args: - stats: Column statistics from pg_stats - is_lower: True if we want the lower bound, False for upper bound - - Returns: - Appropriate tight bound value based on available statistics - """ - data_type = stats.get("data_type", "").lower() - - # First check for most common values - these are statistically most relevant - common_vals = stats.get("common_vals") - common_freqs = stats.get("common_freqs") - - if common_vals and common_freqs and len(common_vals) == len(common_freqs): - # Use the most common value if available - if len(common_vals) > 0: - common_vals_list = list(common_vals) # make sure it's a list - common_freqs_list = list(common_freqs) # make sure it's a list - # Find the most frequent value - max_freq_idx = common_freqs_list.index(max(common_freqs_list)) - most_common = common_vals_list[max_freq_idx] - - # For tight bounds, use the most common value with small adjustment - try: - if isinstance(most_common, float): - # Small +/- adjustment around most common value - adjustment = abs(most_common) * 0.05 if most_common != 0 else 1 - return most_common - adjustment if is_lower else most_common + adjustment - if isinstance(most_common, int): - # Small +/- adjustment around most common value - adjustment = abs(most_common) * 0.05 if most_common != 0 else 1 - return int(most_common - adjustment) if is_lower else int(most_common + adjustment) - elif isinstance(most_common, str) and most_common.isdigit(): - # For string digits, convert and adjust - num_val = float(most_common) - adjustment = abs(num_val) * 0.05 if num_val != 0 else 1 - return str(int(num_val - adjustment)) if is_lower else str(int(num_val + adjustment)) - else: - # For non-numeric, just use most common - return most_common - except (TypeError, ValueError): - logger.warning(f"Error adapting most common value: {most_common}") - # If adaptation fails, just use the value - return most_common - - # Next, try histogram bounds focusing on central values - histogram_bounds = stats.get("histogram_bounds") - if histogram_bounds and len(histogram_bounds) >= 3: - # For tight bounds, use values very close to median - median_idx = len(histogram_bounds) // 2 - # Only move 10% away from median in either direction for tight bounds - idx_offset = max(1, len(histogram_bounds) // 10) - - if is_lower: - bound_idx = max(0, median_idx - idx_offset) - else: - bound_idx = min(len(histogram_bounds) - 1, median_idx + idx_offset) - - return histogram_bounds[bound_idx] - - # Fall back to standard statistics if available - most_common = stats.get("most_common_vals", [None])[0] if stats.get("most_common_vals") else None - if most_common is not None: - return most_common - - # Use very conservative defaults as last resort - if "int" in data_type or data_type in ["smallint", "integer", "bigint"]: - return 10 if is_lower else 20 # Very tight range - elif data_type in ["numeric", "decimal", "real", "double precision", "float"]: - return 10.0 if is_lower else 20.0 # Very tight range - elif "date" in data_type or "time" in data_type: - return "'2023-01-01'" if is_lower else "'2023-01-31'" # Just one month - elif data_type == "boolean": - return "true" # Same value for both bounds for boolean - else: - # Default string-like behavior - narrow range - return "'m'" if is_lower else "'n'" # Just two adjacent letters - - def _get_table_aliases(self, query: str, table_name: str) -> list[str]: - """Extract table aliases for a given table using SQL parser. - - Args: - query: The SQL query to parse - table_name: The name of the table to find aliases for - - Returns: - List of aliases (including the table name itself) - """ - try: - # Parse the query - parsed = parse_sql(query) - if not parsed: - return [table_name] # Return just the table name if parsing fails - - # Get the statement tree - stmt = parsed[0].stmt - - # Use TableAliasVisitor to extract aliases - alias_visitor = TableAliasVisitor() - alias_visitor(stmt) - - # Find all aliases for this table - aliases = [table_name] # Always include the table name itself - - for alias, table in alias_visitor.aliases.items(): - if table.lower() == table_name.lower(): - aliases.append(alias) - - return aliases - except Exception as e: - logger.error(f"Error extracting table aliases: {e}", exc_info=True) - return [table_name] # Fallback to just the table name - - def _identify_parameter_column(self, context: str, table_columns: dict[str, set[str]]) -> tuple[str, str] | None: - """Identify which column a parameter likely belongs to based on context.""" - # Look for patterns like "column_name = $1" or "column_name IN ($1)" - for table, columns in table_columns.items(): - for column in columns: - # Various patterns to match column references - patterns = [ - rf"{column}\s*=\s*\$\d+", # column = $1 - rf"{column}\s+in\s+\([^)]*\$\d+[^)]*\)", # column in (...$1...) - rf"{column}\s+like\s+\$\d+", # column like $1 - rf"{column}\s*>\s*\$\d+", # column > $1 - rf"{column}\s*<\s*\$\d+", # column < $1 - rf"{column}\s*>=\s*\$\d+", # column >= $1 - rf"{column}\s*<=\s*\$\d+", # column <= $1 - rf"{column}\s+between\s+\$\d+\s+and\s+\$\d+", # column between $1 and $2 - ] - - for pattern in patterns: - if re.search(pattern, context, re.IGNORECASE): - return (table, column) - - return None - - async def _get_column_statistics(self, table_name: str, column_name: str) -> dict[str, Any] | None: - """Get statistics for a column from pg_stats.""" - # Create a cache key from table and column name - cache_key = f"{table_name}.{column_name}" - - # Check if we already have this in cache - if cache_key in self._column_stats_cache: - return self._column_stats_cache[cache_key] - - # Not in cache, query the database - try: - query = """ - SELECT - data_type, - most_common_vals as common_vals, - most_common_freqs as common_freqs, - histogram_bounds, - null_frac, - n_distinct, - correlation - FROM pg_stats - JOIN information_schema.columns - ON pg_stats.tablename = information_schema.columns.table_name - AND pg_stats.attname = information_schema.columns.column_name - WHERE pg_stats.tablename = {} - AND pg_stats.attname = {} - """ - - result = await SafeSqlDriver.execute_param_query( - self.sql_driver, - query, - [table_name, column_name], - ) - if not result or not result[0]: - self._column_stats_cache[cache_key] = None - return None - - stats = dict(result[0].cells) - - # Convert PostgreSQL arrays to Python lists for easier handling - for key in ["common_vals", "common_freqs", "histogram_bounds"]: - if key in stats and stats[key] is not None: - if isinstance(stats[key], str): - # Parse array literals like '{val1,val2}' into Python lists - array_str = stats[key].strip("{}") - if array_str: - stats[key] = [self._parse_pg_array_value(val) for val in array_str.split(",")] - else: - stats[key] = [] - - # Cache the processed results - self._column_stats_cache[cache_key] = stats - return stats - except Exception as e: - logger.warning(f"Error getting column statistics for {table_name}.{column_name}: {e}") - self._column_stats_cache[cache_key] = None - return None - - def _parse_pg_array_value(self, value: str) -> Any: - """Parse a single value from a PostgreSQL array representation.""" - value = value.strip() - - # Try to convert to appropriate type - if value == "null": - return None - elif value.startswith('"') and value.endswith('"'): - return value[1:-1] # Strip quotes for string values - - # Try numeric conversion - try: - if "." in value: - return float(value) - else: - return int(value) - except ValueError: - # Return as string if not a number - return value - - def _get_replacement_value(self, stats: dict[str, Any], context: str) -> str: - """Generate an appropriate replacement value based on column statistics.""" - data_type = stats.get("data_type", "").lower() - common_vals = stats.get("common_vals") - histogram_bounds = stats.get("histogram_bounds") - - # If we have common values, use the most common one for equality, - # or a value in the middle of the range for range queries - - # Detect query operator context - is_equality = "=" in context and "!=" not in context and "<>" not in context - is_range = any(op in context for op in [">", "<", ">=", "<=", "between"]) - is_like = "like" in context - - # For string types - if "char" in data_type or data_type == "text": - if is_like: - return "'%test%'" - elif common_vals and is_equality: - # Use the most common value for equality - sample = common_vals[0] - return f"'{sample}'" - elif common_vals: - # Use any sample value - sample = common_vals[0] - return f"'{sample}'" - else: - # Default string - return "'sample_value'" - - # For numeric types - elif "int" in data_type or data_type in [ - "numeric", - "decimal", - "real", - "double", - ]: - if histogram_bounds and is_range: - # For range queries, use a value in the middle - bounds = histogram_bounds - if isinstance(bounds, list) and len(bounds) > 1: - middle_idx = len(bounds) // 2 - value = bounds[middle_idx] - return str(value) - elif common_vals and is_equality: - # Use most common value for equality - return str(common_vals[0]) - elif histogram_bounds: - # Use a reasonable value from histogram - bounds = histogram_bounds - if isinstance(bounds, list) and len(bounds) > 0: - return str(bounds[0]) - - # Default numeric values by type - if "int" in data_type: - return "41" - else: - return "41.5" - - # For date/time types - elif "date" in data_type or "time" in data_type: - if is_range: - return "'2023-01-15'" # Middle of the month - return "'2023-01-01'" - - # For boolean - elif data_type == "boolean": - return "true" - - # Default fallback - return "'sample_value'" - - def _get_generic_replacement(self, context: str) -> str: - """Provide a generic replacement when we can't determine the specific column type.""" - context = context.lower() - - # Try to guess based on context - if any(date_word in context.split() for date_word in ["date", "timestamp", "time"]): - return "'2023-01-01'" - - if any(word in context for word in ["id", "key", "code", "num"]): - return "43" - - if "like" in context: - return "'%sample%'" - - if any(word in context for word in ["amount", "price", "cost", "fee"]): - return "99.99" - - # If in comparison context, use a number - if any(op in context for op in ["=", ">", "<", ">=", "<="]): - return "44" - - # Default to string - return "'sample_value'" - - def _replace_parameters_generic(self, query: str) -> str: - """Fallback generic parameter replacement when catalog lookup fails.""" - try: - modified_query = query - - # Replace string parameters - modified_query = re.sub(r"like \$\d+", "like '%'", modified_query) - - # Context-aware replacements - modified_query = re.sub( - r"(\w+)\s*=\s*\$\d+", - lambda m: self._context_replace(m, "="), - modified_query, - ) - modified_query = re.sub( - r"(\w+)\s*<\s*\$\d+", - lambda m: self._context_replace(m, "<"), - modified_query, - ) - modified_query = re.sub( - r"(\w+)\s*>\s*\$\d+", - lambda m: self._context_replace(m, ">"), - modified_query, - ) - - # Replace numeric parameters in inequalities - modified_query = re.sub(r"(\d+) and \$\d+", r"\1 and 100", modified_query) - modified_query = re.sub(r"\$\d+ and (\d+)", r"1 and \1", modified_query) - modified_query = re.sub(r">\s*\$\d+", "> 1", modified_query) - modified_query = re.sub(r"<\s*\$\d+", "< 100", modified_query) - modified_query = re.sub(r"=\s*\$\d+\b", "= 45", modified_query) - - # For any remaining parameters, use a generic replacement - modified_query = re.sub(r"\$\d+", "'sample_value'", modified_query) - - return modified_query - except Exception as e: - logger.error(f"Error in generic parameter replacement: {e}", exc_info=True) - return query - - def _context_replace(self, match, op: str): - """Replace parameters based on column name context.""" - col_name = match.group(1).lower() - - # ID-like columns (numeric) - if col_name.endswith("id") or col_name.endswith("_id") or col_name == "id": - return f"{col_name} {op} 46" - - # Date/time columns - if any(word in col_name for word in ["date", "time", "created", "updated"]): - return f"{col_name} {op} '2023-01-01'" - - # Numeric-looking columns - if any(word in col_name for word in ["amount", "price", "cost", "count", "num", "qty"]): - return f"{col_name} {op} 46.5" - - # Status-like columns (likely string enums) - if "status" in col_name or "type" in col_name or "state" in col_name: - return f"{col_name} {op} 'active'" - - # Default to string for other columns - return f"{col_name} {op} 'sample_value'" - - def extract_columns(self, query: str) -> dict[str, set[str]]: - """Extract columns from a query using improved visitors.""" - try: - parsed = parse_sql(query) - if not parsed: - return {} - stmt = parsed[0].stmt - if not isinstance(stmt, SelectStmt): - return {} - - return self.extract_stmt_columns(stmt) - - except Exception: - logger.warning(f"Error extracting columns from query: {query}") - return {} - - def extract_stmt_columns(self, stmt: SelectStmt) -> dict[str, set[str]]: - """Extract columns from a query using improved visitors.""" - try: - # Second pass: collect columns with table context - collector = ColumnCollector() - collector(stmt) - - return collector.columns - - except Exception: - logger.warning(f"Error extracting columns from query: {stmt}") - return {} diff --git a/src/mcpg/_vendor/sql/extension_utils.py b/src/mcpg/_vendor/sql/extension_utils.py deleted file mode 100644 index 1f96fa4..0000000 --- a/src/mcpg/_vendor/sql/extension_utils.py +++ /dev/null @@ -1,246 +0,0 @@ -"""Utilities for working with PostgreSQL extensions.""" - -import logging -from dataclasses import dataclass -from typing import Literal - -from .safe_sql import SafeSqlDriver -from .sql_driver import SqlDriver - -logger = logging.getLogger(__name__) - -# Single global PostgreSQL version cache -# TODO: If we support multiple connections in the future, this should be connection-specific -_POSTGRES_VERSION = None - - -@dataclass -class ExtensionStatus: - """Status of an extension.""" - - is_installed: bool - is_available: bool - name: str - message: str - default_version: str | None - - -def reset_postgres_version_cache() -> None: - """Reset the PostgreSQL version cache. Primarily used for testing.""" - global _POSTGRES_VERSION - _POSTGRES_VERSION = None - - -async def get_postgres_version(sql_driver: SqlDriver) -> int: - """ - Get the major PostgreSQL version as an integer. - - Args: - sql_driver: An instance of SqlDriver to execute queries - - Returns: - The major PostgreSQL version as an integer (e.g., 16 for PostgreSQL 16.2) - Returns 0 if the version cannot be determined - """ - # Check if we have a cached version - global _POSTGRES_VERSION - if _POSTGRES_VERSION is not None: - return _POSTGRES_VERSION - - try: - rows = await sql_driver.execute_query("SHOW server_version") - if not rows: - logger.warning("Could not determine PostgreSQL version") - return 0 - - version_string = rows[0].cells["server_version"] - # Extract the major version (before the first dot) - major_version = version_string.split(".")[0] - version = int(major_version) - - # Cache the version globally - _POSTGRES_VERSION = version - - return version - except Exception as e: - raise ValueError("Error determining PostgreSQL version") from e - - -async def check_postgres_version_requirement(sql_driver: SqlDriver, min_version: int, feature_name: str) -> tuple[bool, str]: - """ - Check if the PostgreSQL version meets the minimum requirement. - - Args: - sql_driver: An instance of SqlDriver to execute queries - min_version: The minimum required PostgreSQL version - feature_name: Name of the feature that requires this version - - Returns: - A tuple of (meets_requirement, message) - """ - pg_version = await get_postgres_version(sql_driver) - - if pg_version >= min_version: - return True, f"PostgreSQL version {pg_version} meets the requirement for {feature_name}" - - return False, ( - f"This feature ({feature_name}) requires PostgreSQL {min_version} or later. Your current version is PostgreSQL {pg_version or 'unknown'}." - ) - - -async def check_extension( - sql_driver: SqlDriver, - extension_name: str, - include_messages: bool = True, - message_type: Literal["plain", "markdown"] = "plain", -) -> ExtensionStatus: - """ - Check if a PostgreSQL extension is installed or available. - - Args: - sql_driver: An instance of SqlDriver to execute queries - extension_name: Name of the extension to check - include_messages: Whether to include user-friendly messages in the result - message_type: Format for messages - 'plain' or 'markdown' - - Returns: - ExtensionStatus with fields: - - is_installed: True if the extension is installed - - is_available: True if the extension is available but not installed - - name: The extension name - - message: A user-friendly message about the extension status - - default_version: The default version of the extension if available - """ - # Check if the extension is installed - installed_result = await SafeSqlDriver.execute_param_query( - sql_driver, - "SELECT extversion FROM pg_extension WHERE extname = {}", - [extension_name], - ) - - # Initialize result - result = ExtensionStatus( - is_installed=False, - is_available=False, - name=extension_name, - message="", - default_version=None, - ) - - if installed_result and len(installed_result) > 0: - # Extension is installed - version = installed_result[0].cells.get("extversion", "unknown") - result.is_installed = True - result.is_available = True - - if include_messages: - if message_type == "markdown": - result.message = f"The **{extension_name}** extension (version {version}) is already installed." - else: - result.message = f"The {extension_name} extension (version {version}) is already installed." - else: - # Check if the extension is available but not installed - available_result = await SafeSqlDriver.execute_param_query( - sql_driver, - "SELECT default_version FROM pg_available_extensions WHERE name = {}", - [extension_name], - ) - - if available_result and len(available_result) > 0: - # Extension is available but not installed - result.is_available = True - result.default_version = available_result[0].cells.get("default_version") - - if include_messages: - if message_type == "markdown": - result.message = ( - f"The **{extension_name}** extension is available but not installed.\n\n" - f"You can install it by running: `CREATE EXTENSION {extension_name};`." - ) - else: - result.message = ( - f"The {extension_name} extension is available but not installed.\n" - f"You can install it by running: CREATE EXTENSION {extension_name};" - ) - else: - # Extension is not available - if include_messages: - if message_type == "markdown": - result.message = ( - f"The **{extension_name}** extension is not available on this PostgreSQL server.\n\n" - f"To install it, you need to:\n" - f"1. Install the extension package on the server\n" - f"2. Run: `CREATE EXTENSION {extension_name};`" - ) - else: - result.message = ( - f"The {extension_name} extension is not available on this PostgreSQL server.\n" - f"To install it, you need to:\n" - f"1. Install the extension package on the server\n" - f"2. Run: CREATE EXTENSION {extension_name};" - ) - - return result - - -async def check_hypopg_installation_status(sql_driver: SqlDriver, message_type: Literal["plain", "markdown"] = "markdown") -> tuple[bool, str]: - """ - Get a detailed status message for the HypoPG extension. - - Args: - sql_driver: An instance of SqlDriver to execute queries - message_type: Format for messages - 'plain' or 'markdown' - - Returns: - A formatted message about the HypoPG extension status with installation instructions - """ - status = await check_extension(sql_driver, "hypopg", include_messages=False) - - if status.is_installed: - if message_type == "markdown": - return True, "The **hypopg** extension is already installed." - else: - return True, "The hypopg extension is already installed." - - if status.is_available: - if message_type == "markdown": - return False, ( - "The **hypopg** extension is required to test hypothetical indexes, but it is not currently installed.\n\n" - "You can ask me to install 'hypopg' using the 'execute_query' tool.\n\n" - "**Is it safe?** Installing 'hypopg' is generally safe and a standard practice for index testing. " - "It adds a virtual layer that simulates indexes without actually creating them in the database. " - "It requires database privileges (often superuser) to install.\n\n" - "**What does it do?** It allows you to create virtual indexes and test how they would affect query performance " - "without the overhead of actually creating the indexes.\n\n" - "**How to undo?** If you later decide to remove it, you can ask me to run 'DROP EXTENSION hypopg;'." - ) - else: - return False, ( - "The hypopg extension is required to test hypothetical indexes, but it is not currently installed.\n" - "You can ask me to install it using the 'execute_query' tool.\n" - "It is generally safe to install and allows testing indexes without creating them." - ) - - pg_version = await get_postgres_version(sql_driver) - major_version_str = f"{pg_version}" if pg_version > 0 else "XX" - # Extension is not available - if message_type == "markdown": - return False, ( - "The **hypopg** extension is not available on this PostgreSQL server.\n\n" - "To install HypoPG:\n" - f"1. For Debian/Ubuntu: `sudo apt-get install postgresql-{major_version_str}-hypopg`\n" - f"2. For RHEL/CentOS: `sudo yum install postgresql{major_version_str}-hypopg`\n" - "3. For MacOS with Homebrew: `brew install hypopg`\n" - "4. For other systems, build from source: `git clone https://github.com/HypoPG/hypopg`\n\n" - "After installing the extension packages, connect to your database and run: `CREATE EXTENSION hypopg;`" - ) - else: - return False, ( - "The hypopg extension is not available on this PostgreSQL server.\n" - "To install HypoPG:\n" - f"1. For Debian/Ubuntu: sudo apt-get install postgresql-{major_version_str}-hypopg\n" - f"2. For RHEL/CentOS: sudo yum install postgresql{major_version_str}-hypopg\n" - "3. For MacOS with Homebrew: brew install hypopg\n" - "4. For other systems, build from source: git clone https://github.com/HypoPG/hypopg\n" - "After installing the extension packages, connect to your database and run: CREATE EXTENSION hypopg;" - ) diff --git a/src/mcpg/_vendor/sql/index.py b/src/mcpg/_vendor/sql/index.py deleted file mode 100644 index e1564b5..0000000 --- a/src/mcpg/_vendor/sql/index.py +++ /dev/null @@ -1,52 +0,0 @@ -from dataclasses import dataclass -from typing import Any - - -@dataclass(frozen=True) -class IndexDefinition: - """Immutable index configuration for hashing.""" - - table: str - columns: tuple[str, ...] - using: str = "btree" - - def to_dict(self) -> dict[str, Any]: - return { - "table": self.table, - "columns": list(self.columns), - "using": self.using, - "definition": self.definition, - } - - @property - def definition(self) -> str: - return f"CREATE INDEX {self.name} ON {self.table} USING {self.using} ({', '.join(self.columns)})" - - @property - def name(self) -> str: - # Clean column names for use in index naming - # Replace special characters with underscores to avoid issues with - # functional expressions - cleaned_columns = [] - for col in self.columns: - # Replace parentheses and other special characters with underscores - # This ensures expressions like LOWER(column_name) work in - # index names - cleaned_col = col.replace("(", "_").replace(")", "_").replace(" ", "_").replace(",", "_") - # Remove consecutive underscores - while "__" in cleaned_col: - cleaned_col = cleaned_col.replace("__", "_") - # Remove trailing underscores - cleaned_col = cleaned_col.rstrip("_") - cleaned_columns.append(cleaned_col) - - column_part = "_".join(cleaned_columns) - suffix = "" if self.using == "btree" else f"_{self.using}" - base = f"crystaldba_idx_{self.table}_{column_part}_{len(self.columns)}" - return f"{base}{suffix}" - - def __str__(self) -> str: - return self.definition - - def __repr__(self) -> str: - return f"IndexConfig(table='{self.table}', columns={self.columns}, using='{self.using}')" diff --git a/src/mcpg/_vendor/sql/safe_sql.py b/src/mcpg/_vendor/sql/safe_sql.py deleted file mode 100644 index 37382f0..0000000 --- a/src/mcpg/_vendor/sql/safe_sql.py +++ /dev/null @@ -1,1036 +0,0 @@ -from __future__ import annotations - -import asyncio -import logging -import re -from typing import Any -from typing import ClassVar -from typing import Optional - -import pglast -from pglast.ast import A_ArrayExpr -from pglast.ast import A_Const -from pglast.ast import A_Expr -from pglast.ast import A_Indices -from pglast.ast import A_Indirection -from pglast.ast import A_Star -from pglast.ast import Alias -from pglast.ast import BitString -from pglast.ast import Boolean -from pglast.ast import BooleanTest -from pglast.ast import BoolExpr -from pglast.ast import CaseExpr -from pglast.ast import CaseWhen -from pglast.ast import ClosePortalStmt -from pglast.ast import CoalesceExpr -from pglast.ast import CollateClause -from pglast.ast import ColumnRef -from pglast.ast import CommonTableExpr -from pglast.ast import CreateExtensionStmt -from pglast.ast import DeallocateStmt -from pglast.ast import DeclareCursorStmt -from pglast.ast import DefElem -from pglast.ast import ExplainStmt -from pglast.ast import FetchStmt -from pglast.ast import Float -from pglast.ast import FromExpr -from pglast.ast import FuncCall -from pglast.ast import GroupingFunc -from pglast.ast import GroupingSet -from pglast.ast import Integer -from pglast.ast import JoinExpr -from pglast.ast import MinMaxExpr -from pglast.ast import NamedArgExpr -from pglast.ast import Node -from pglast.ast import NotifyStmt -from pglast.ast import NullTest -from pglast.ast import ParamRef -from pglast.ast import PrepareStmt -from pglast.ast import RangeFunction -from pglast.ast import RangeSubselect -from pglast.ast import RangeTableFunc -from pglast.ast import RangeTableFuncCol -from pglast.ast import RangeTableSample -from pglast.ast import RangeVar -from pglast.ast import RawStmt -from pglast.ast import ResTarget -from pglast.ast import RowCompareExpr -from pglast.ast import RowExpr -from pglast.ast import ScalarArrayOpExpr -from pglast.ast import SelectStmt -from pglast.ast import SortBy -from pglast.ast import SortGroupClause -from pglast.ast import SQLValueFunction -from pglast.ast import String -from pglast.ast import SubLink -from pglast.ast import TableFunc -from pglast.ast import TableSampleClause -from pglast.ast import TargetEntry -from pglast.ast import TypeCast -from pglast.ast import TypeName -from pglast.ast import VacuumStmt -from pglast.ast import VariableShowStmt -from pglast.ast import WindowClause -from pglast.ast import WindowDef -from pglast.ast import WindowFunc -from pglast.ast import WithClause -from pglast.enums import A_Expr_Kind -from psycopg.sql import SQL -from psycopg.sql import Composable -from psycopg.sql import Literal -from typing_extensions import LiteralString - -from .sql_driver import SqlDriver - -logger = logging.getLogger(__name__) - - -class SafeSqlDriver(SqlDriver): - """A wrapper around any SqlDriver that only allows SELECT, ANALYZE, VACUUM, EXPLAIN SELECT, and - SHOW queries. - - Uses pglast to parse and validate SQL statements before execution. - All other statement types (DDL, DML etc) are rejected. - Performs deep validation of the query tree to prevent unsafe operations. - """ - - # Pattern to match pg_catalog schema qualification - PG_CATALOG_PATTERN = re.compile(r"^pg_catalog\.(.+)$") - # Pattern to validate LIKE expressions - must either start with % or end with %, but not both - LIKE_PATTERN = re.compile(r"^[^%]+%$") - - ALLOWED_STMT_TYPES: ClassVar[set[type]] = { - SelectStmt, # Regular SELECT - ExplainStmt, # EXPLAIN SELECT - CreateExtensionStmt, # CREATE EXTENSION - VariableShowStmt, # SHOW statements - VacuumStmt, # VACUUM and ANALYZE statements - PrepareStmt, # PREPARE statement (for prepared queries) - # ExecuteStmt, # EXECUTE statement (for prepared queries) - DeallocateStmt, # DEALLOCATE statement (for prepared queries) - # CopyStmt, # COPY TO (for exporting query results - will be validated to ensure COPY FROM is not allowed) - # DiscardStmt, # DISCARD (for discarding client-side caches) - DeclareCursorStmt, # DECLARE CURSOR (for cursor operations) - ClosePortalStmt, # CLOSE (for closing cursors) - FetchStmt, # FETCH (for retrieving cursor results) - # CheckPointStmt, # CHECKPOINT (read-only administrative command) - # ListenStmt, # LISTEN (notification system) - # UnlistenStmt, # UNLISTEN (notification system) - # VariableSetStmt, # SET (for session variables - will be validated) - } - - ALLOWED_FUNCTIONS: ClassVar[set[str]] = { - # Aggregate functions - "array_agg", - "avg", - "bit_and", - "bit_or", - "bool_and", - "bool_or", - "count", - "every", - "json_agg", - "jsonb_agg", - "json_object_agg", - "jsonb_object_agg", - "max", - "min", - "string_agg", - "sum", - "xmlagg", - # Mathematical functions - "abs", - "cbrt", - "ceil", - "ceiling", - "degrees", - "div", - "erf", - "erfc", - "exp", - "factorial", - "floor", - "gcd", - "lcm", - "ln", - "log", - "log10", - "min_scale", - "mod", - "pi", - "power", - "radians", - "random", - "random_normal", - "round", - "scale", - "setseed", - "sign", - "sqrt", - "trim_scale", - "trunc", - "width_bucket", - "percentile_cont", - "percentile_disc", - "mode", - # Trigonometric functions - "acos", - "acosd", - "asin", - "asind", - "atan", - "atand", - "atan2", - "atan2d", - "cos", - "cosd", - "cot", - "cotd", - "sin", - "sind", - "tan", - "tand", - # Hyperbolic functions - "sinh", - "cosh", - "tanh", - "asinh", - "acosh", - "atanh", - # Array functions and operators - "array", - "array_append", - "array_cat", - "array_dims", - "array_fill", - "array_length", - "array_lower", - "array_ndims", - "array_position", - "array_positions", - "array_prepend", - "array_remove", - "array_replace", - "array_sample", - "array_shuffle", - "array_to_string", - "array_upper", - "cardinality", - "string_to_array", - "trim_array", - "unnest", - "any", - # String functions - "ascii", - "bit_length", - "btrim", - "char_length", - "character_length", - "chr", - "concat", - "concat_ws", - "convert", - "convert_from", - "convert_to", - "decode", - "encode", - "format", - "initcap", - "left", - "length", - "lower", - "lpad", - "ltrim", - "md5", - "normalize", - "octet_length", - "overlay", - "parse_ident", - "position", - "quote_ident", - "quote_literal", - "quote_nullable", - "repeat", - "replace", - "reverse", - "right", - "rpad", - "rtrim", - "split_part", - "starts_with", - "string_to_table", - "strpos", - "substr", - "substring", - "to_ascii", - "to_bin", - "to_hex", - "to_oct", - "translate", - "trim", - "upper", - "unistr", - "unicode_assigned", - # Pattern matching functions - "regexp_match", - "regexp_matches", - "regexp_replace", - "regexp_split_to_array", - "regexp_split_to_table", - "regexp_substr", - "regexp_count", - "regexp_instr", - "regexp_like", - # Type casting functions - "regclass", - "to_char", - "to_date", - "to_number", - "to_timestamp", - # Information functions - "current_catalog", - "current_database", - "current_query", - "current_role", - "current_schema", - "current_schemas", - "current_setting", - "current_user", - "pg_backend_pid", - "pg_blocking_pids", - "pg_conf_load_time", - "pg_current_logfile", - "pg_jit_available", - "pg_safe_snapshot_blocking_pids", - "pg_trigger_depth", - "session_user", - "system_user", - "user", - "version", - "unicode_version", - "icu_unicode_version", - # Database object information functions - "pg_column_size", - "pg_column_compression", - "pg_column_toast_chunk_id", - "pg_database_size", - "pg_indexes_size", - "pg_get_indexdef", - "pg_relation_filenode", - "pg_relation_size", - "pg_size_bytes", - "pg_size_pretty", - "pg_table_size", - "pg_tablespace_size", - "pg_total_relation_size", - # Security privilege check functions - "has_any_column_privilege", - "has_column_privilege", - "has_database_privilege", - "has_foreign_data_wrapper_privilege", - "has_function_privilege", - "has_language_privilege", - "has_parameter_privilege", - "has_schema_privilege", - "has_sequence_privilege", - "has_server_privilege", - "has_table_privilege", - "has_tablespace_privilege", - "has_type_privilege", - "pg_has_role", - "row_security_active", - # JSON functions - "json", - "json_array_length", - "jsonb_array_length", - "json_each", - "jsonb_each", - "json_each_text", - "jsonb_each_text", - "json_extract_path", - "jsonb_extract_path", - "json_extract_path_text", - "jsonb_extract_path_text", - "json_object_keys", - "jsonb_object_keys", - "json_array_elements", - "jsonb_array_elements", - "json_array_elements_text", - "jsonb_array_elements_text", - "json_typeof", - "jsonb_typeof", - "json_strip_nulls", - "jsonb_strip_nulls", - "jsonb_set", - "jsonb_set_path", - "jsonb_set_lax", - "jsonb_pretty", - "json_build_array", - "jsonb_build_array", - "json_build_object", - "jsonb_build_object", - "json_object", - "jsonb_object", - "json_scalar", - "json_serialize", - "json_populate_record", - "jsonb_populate_record", - "jsonb_populate_record_valid", - "json_populate_recordset", - "jsonb_populate_recordset", - "json_to_record", - "jsonb_to_record", - "json_to_recordset", - "jsonb_to_recordset", - "jsonb_insert", - "jsonb_path_exists", - "jsonb_path_match", - "jsonb_path_query", - "jsonb_path_query_array", - "jsonb_path_query_first", - "jsonb_path_exists_tz", - "jsonb_path_match_tz", - "jsonb_path_query_tz", - "jsonb_path_query_array_tz", - "jsonb_path_query_first_tz", - "jsonbv_typeof", - "to_json", - "to_jsonb", - "array_to_json", - "row_to_json", - # Object Info - "pg_get_expr", - "pg_get_functiondef", - "pg_get_function_arguments", - "pg_get_function_identity_arguments", - "pg_get_function_result", - "pg_get_catalog_foreign_keys", - "pg_get_constraintdef", - "pg_get_userbyid", - "pg_get_keywords", - "pg_get_partkeydef", - # Encoding Functions - "pg_basetype", - "pg_client_encoding", - "pg_encoding_to_char", - "pg_char_to_encoding", - # Validity Checking - "pg_input_is_valid", - "pg_input_error_info", - # Object Definition/Information Functions - "pg_get_serial_sequence", - "pg_get_viewdef", - "pg_get_ruledef", - "pg_get_triggerdef", - "pg_get_statisticsobjdef", - # Type Information and Conversion - "pg_typeof", - "format_type", - "to_regtype", - "to_regtypemod", - # Object Name/OID Translation - "to_regclass", - "to_regcollation", - "to_regnamespace", - "to_regoper", - "to_regoperator", - "to_regproc", - "to_regprocedure", - "to_regrole", - # Index Property Functions - "pg_index_column_has_property", - "pg_index_has_property", - "pg_indexam_has_property", - # Schema Visibility Functions - "pg_collation_is_visible", - "pg_conversion_is_visible", - "pg_function_is_visible", - "pg_opclass_is_visible", - "pg_operator_is_visible", - "pg_opfamily_is_visible", - "pg_statistics_obj_is_visible", - "pg_table_is_visible", - "pg_ts_config_is_visible", - "pg_ts_dict_is_visible", - "pg_ts_parser_is_visible", - "pg_ts_template_is_visible", - "pg_type_is_visible", - # Date/Time Functions - "age", # When used with timestamp arguments (not transaction IDs) - "clock_timestamp", - "current_date", - "current_time", - "current_timestamp", - "date_part", - "date_trunc", - "extract", - "isfinite", - "justify_days", - "justify_hours", - "justify_interval", - "localtime", - "localtimestamp", - "make_date", - "make_interval", - "make_time", - "make_timestamp", - "make_timestamptz", - "now", - "statement_timestamp", - "timeofday", - "transaction_timestamp", - # Additional Type Conversion - "cast", - "text", - "bool", - "int2", - "int4", - "int8", - "float4", - "float8", - "numeric", - "date", - "time", - "timetz", - "timestamp", - "timestamptz", - "interval", - # ACL functions - "acldefault", - "aclexplode", - "makeaclitem", - # System/Configuration Information - "pg_tablespace_location", # Exposes filesystem paths - "pg_tablespace_databases", # Exposes system-wide information - "pg_settings_get_flags", # Exposes configuration details - "pg_options_to_table", # Exposes internal storage options - # Transaction/WAL Information - "pg_current_xact_id", # Exposes transaction details - "pg_current_snapshot", # Exposes transaction snapshots - "pg_snapshot_xip", # Exposes in-progress transactions - "pg_xact_commit_timestamp", # Transaction timestamp details - # Control Data Functions - "pg_control_checkpoint", # Internal checkpoint state - "pg_control_system", # Control file state - "pg_control_init", # Cluster initialization state - "pg_control_recovery", # Recovery state - # WAL Functions - "pg_available_wal_summaries", # WAL summary information - "pg_wal_summary_contents", # WAL contents - "pg_get_wal_summarizer_state", # WAL summarizer state - # Network Information - "inet_client_addr", # Client network details - "inet_client_port", # Client port information - "inet_server_addr", # Server network details - "inet_server_port", # Server port information - # Temporary Schema Information - "pg_my_temp_schema", # Exposes temp schema OIDs - "pg_is_other_temp_schema", # Shows other sessions' temp schemas - # Notification/Channel Information - "pg_listening_channels", # Shows active notification channels - "pg_notification_queue_usage", # Exposes notification queue state - # Server State Information - "pg_postmaster_start_time", # Shows server start time - # Recovery Information Functions (safe ones) - "pg_is_in_recovery", - # Hypopg functions - "hypopg_create_index", - "hypopg_reset", - "hypopg_relation_size", - "hypopg_list_indexes", - "hypopg_get_indexdef", - "hypopg_hide_index", - "hypopg_unhide_index", - # XML Functions (read-only) - "xml", - "xmlcomment", - "xmlconcat", - "xmlelement", - "xmlforest", - "xmlpi", - "xmlroot", - "xmlexists", - "xml_is_well_formed", - "xml_is_well_formed_document", - "xml_is_well_formed_content", - "xpath", - "xpath_exists", - "xmltable", - "xmlnamespaces", - # Network/IP Address Functions - "abbrev", - "broadcast", - "cidr_split", - "family", - "host", - "hostmask", - "inet_merge", - "inet_same_family", - "macaddr8_set7bit", - "masklen", - "netmask", - "network", - "set_bit", - "set_masklen", - # Full-Text Search Functions - "array_to_tsvector", - "get_current_ts_config", - "numnode", - "plainto_tsquery", - "phraseto_tsquery", - "querytree", - "setweight", - "strip", - "to_tsquery", - "to_tsvector", - "ts_debug", - "ts_delete", - "ts_filter", - "ts_headline", - "ts_lexize", - "ts_parse", - "ts_rank", - "ts_rank_cd", - "ts_rewrite", - "ts_stat", - "ts_token_type", - "tsvector_to_array", - "websearch_to_tsquery", - # Range Functions - "isempty", - "lower_inc", - "upper_inc", - "lower_inf", - "upper_inf", - "range_merge", - # Geometry Functions - "area", - "center", - "diameter", - "height", - "isclosed", - "isopen", - "npoints", - "pclose", - "popen", - "radius", - "width", - # UUID Functions - "gen_random_uuid", - # Enum Functions - "enum_first", - "enum_last", - "enum_range", - # Sequence Functions (read-only) - "currval", - "lastval", - # pg_trgm Functions - "similarity", - "show_trgm", - "show_limit", - "word_similarity", - "strict_word_similarity", - # pg_crypto Functions (read-only) - "digest", - "hmac", - "encrypt", - "gen_salt", - "crypt", - # earthdistance/cube Functions - "earth_distance", - "earth_box", - "ll_to_earth", - "earth_to_ll", - "cube_distance", - # fuzzystrmatch Functions - "soundex", - "difference", - "levenshtein", - "metaphone", - "dmetaphone", - # Window functions - "row_number", - "rank", - "dense_rank", - "percent_rank", - "cume_dist", - "ntile", - "lag", - "lead", - "first_value", - "last_value", - "nth_value", - } - - ALLOWED_NODE_TYPES: ClassVar[set[type]] = ALLOWED_STMT_TYPES | { - # Basic SELECT components - ResTarget, - ColumnRef, - A_Star, - A_Const, - A_Expr, - BoolExpr, - BooleanTest, - NullTest, - RangeVar, - JoinExpr, - FromExpr, - WithClause, - CommonTableExpr, - # Basic operators - A_Expr, - BoolExpr, - SubLink, - MinMaxExpr, - RowExpr, - # EXPLAIN components - ExplainStmt, - DefElem, - # SHOW components - VariableShowStmt, - # VACUUM and ANALYZE components - VacuumStmt, - # Sorting/grouping - SortBy, - SortGroupClause, - # Constants and basic types - Integer, - Float, - String, - BitString, - Boolean, - RawStmt, - ParamRef, - SQLValueFunction, - # Function calls and type casting - FuncCall, - TypeCast, - DefElem, - TypeName, - Alias, - CaseExpr, - CaseWhen, - RangeSubselect, - CoalesceExpr, - NamedArgExpr, - RangeFunction, - A_ArrayExpr, - # Window functions - WindowFunc, - WindowDef, - WindowClause, - # Table functions - TableFunc, - RangeTableFunc, - RangeTableFuncCol, - # Array access - A_Indirection, - A_Indices, - # GROUPING SETS, CUBE, ROLLUP - GroupingSet, - GroupingFunc, - # Table sampling - RangeTableSample, - TableSampleClause, - # Row comparison - RowCompareExpr, - # Collation - CollateClause, - # Others - TargetEntry, - # Additional expression types - ScalarArrayOpExpr, # ARRAY[...] @> or <@ operators - NotifyStmt, # NOTIFY command - } - - ALLOWED_EXTENSIONS: ClassVar[set[str]] = { - # Core PostgreSQL extensions - "hypopg", - "pg_stat_statements", - "pg_trgm", - "btree_gin", - "btree_gist", - "earthdistance", - "cube", - "fuzzystrmatch", - "intarray", - "pgcrypto", - "hstore", - "ltree", - "xml2", - "tablefunc", - "tsm_system_rows", - "tsm_system_time", - "unaccent", - "uuid-ossp", - # Common extensions - "adminpack", - "amcheck", - "bloom", - "citext", - "dict_int", - "dict_xsyn", - "file_fdw", - "intagg", - "isn", - "lo", - "pg_buffercache", - "pg_freespacemap", - "pg_prewarm", - "pg_visibility", - "pgrowlocks", - "pgstattuple", - "plpgsql", - "seg", - "spi", - "sslinfo", - # Foreign data wrappers - "postgres_fdw", - "dblink", - "mysql_fdw", - "mongo_fdw", - "tds_fdw", - "oracle_fdw", - # Popular extensions - "postgis", - "pgrouting", - "timescaledb", - "pg_partman", - "orafce", - "pgaudit", - "pgtap", - "pgsphere", - "pg_qualstats", - "pg_hint_plan", - "auto_explain", - "pg_wait_sampling", - "plv8", - "pg_stat_monitor", - "pg_cron", - "pglogical", - "pgq", - "pgpool_adm", - "pg_fuzzystrmatch", - "pg_bigm", - "pgvector", - "rum", - "zhparser", - "ip4r", - "chkpass", - "tsvector2", - "pg_stat_kcache", - "wal2json", - "pg_repack", - # Programming languages - "plperl", - "plperlu", - "plpython3u", - "plpython", - "pltcl", - "pltclu", - "pljava", - "plrust", - # AWS RDS extensions - "aws_commons", - "aws_s3", - # Supabase extensions - "h3", - "pg_graphql", - "pg_net", - "pgjwt", - "moddatetime", - # Additional data types - "age", - "semver", - "rdkit", - "vector", - # Distributed PostgreSQL - "citus", - # Other validated extensions - "pg_tle", - "pg_roaringbitmap", - "pg_ivm", - "pg_rational", - "pg_partman_bgw", - "q3c", - "pg_track_settings", - "pg_variables", - "pg_walinspect", - "pgmq", - "address_standardizer", - "address_standardizer_data_us", - "postgis_raster", - "postgis_sfcgal", - "postgis_tiger_geocoder", - "postgis_topology", - } - - def __init__(self, sql_driver: SqlDriver, timeout: float | None = None): - """Initialize with an underlying SQL driver and optional timeout. - - Args: - sql_driver: The underlying SQL driver to wrap - timeout: Optional timeout in seconds for query execution - """ - self.sql_driver = sql_driver - self.timeout = timeout - - def _validate_node(self, node: Node) -> None: - """Recursively validate a node and all its children""" - # Check if node type is allowed - if not isinstance(node, tuple(self.ALLOWED_NODE_TYPES)): - raise ValueError(f"Node type {type(node)} is not allowed") - - # Validate LIKE patterns - if isinstance(node, A_Expr) and node.kind in ( - A_Expr_Kind.AEXPR_LIKE, - A_Expr_Kind.AEXPR_ILIKE, - ): - # Get the right-hand side of the LIKE expression (the pattern) - if isinstance(node.rexpr, A_Const) and node.rexpr.val is not None and hasattr(node.rexpr.val, "sval") and node.rexpr.val.sval is not None: - # Nothing to do for now - pass - else: - raise ValueError("LIKE pattern must be a constant string") - - # Validate function calls - if isinstance(node, FuncCall): - func_name = ".".join([str(n.sval) for n in node.funcname]).lower() if node.funcname else "" - # Strip pg_catalog schema if present - match = self.PG_CATALOG_PATTERN.match(func_name) - unqualified_name = match.group(1) if match else func_name - if unqualified_name not in self.ALLOWED_FUNCTIONS: - raise ValueError(f"Function {func_name} is not allowed") - - # Reject SELECT statements with locking clauses - if isinstance(node, SelectStmt) and getattr(node, "lockingClause", None): - raise ValueError("Locking clause on select is prohibited") - - # Reject EXPLAIN ANALYZE statements - if isinstance(node, ExplainStmt): - for option in node.options or []: - if isinstance(option, DefElem) and option.defname == "analyze": - raise ValueError("EXPLAIN ANALYZE is not supported") - - # Reject CREATE EXTENSION statements - if isinstance(node, CreateExtensionStmt): - if node.extname not in self.ALLOWED_EXTENSIONS: - raise ValueError(f"CREATE EXTENSION {node.extname} is not supported") - - # Recursively validate all attributes that might be nodes - for attr_name in node.__slots__: - # Skip private attributes and methods - if attr_name.startswith("_"): - continue - - try: - attr = getattr(node, attr_name) - except AttributeError: - # Skip attributes that don't exist (this is normal in pglast) - continue - - # Handle lists of nodes - if isinstance(attr, list): - for item in attr: - if isinstance(item, Node): - self._validate_node(item) - - # Handle tuples of nodes - elif isinstance(attr, tuple): - for item in attr: - if isinstance(item, Node): - self._validate_node(item) - - # Handle single nodes - elif isinstance(attr, Node): - self._validate_node(attr) - - def _validate(self, query: str) -> None: - """Validate query is safe to execute""" - try: - # Parse the SQL using pglast - parsed = pglast.parse_sql(query) - # Pretty print the parsed SQL for debugging - # print("Parsed SQL:") - # import pprint - # pprint.pprint(parsed) - - # Validate each statement - try: - for stmt in parsed: - if isinstance(stmt, RawStmt): - # Check if the inner statement type is allowed - if not isinstance(stmt.stmt, tuple(self.ALLOWED_STMT_TYPES)): - raise ValueError( - "Only SELECT, ANALYZE, VACUUM, EXPLAIN, SHOW and other read-only statements are allowed. Received raw statement: " - + str(stmt.stmt) - ) - else: - if not isinstance(stmt, tuple(self.ALLOWED_STMT_TYPES)): - raise ValueError( - "Only SELECT, ANALYZE, VACUUM, EXPLAIN, SHOW and other read-only statements are allowed. Received: " + str(stmt) - ) - self._validate_node(stmt) - except Exception as e: - raise ValueError(f"Error validating query: {query}") from e - - except pglast.parser.ParseError as e: - raise ValueError("Failed to parse SQL statement") from e - - async def execute_query( - self, - query: LiteralString, - params: list[Any] | None = None, - force_readonly: bool = True, # do not use value passed in - ) -> Optional[list[SqlDriver.RowResult]]: # noqa: UP007 - """Execute a query after validating it is safe""" - self._validate(query) - - # NOTE: Always force readonly=True in SafeSqlDriver regardless of what was passed - if self.timeout: - try: - async with asyncio.timeout(self.timeout): - return await self.sql_driver.execute_query( - f"/* crystaldba */ {query}", - params=params, - force_readonly=True, - ) - except asyncio.TimeoutError as e: - logger.warning(f"Query execution timed out after {self.timeout} seconds: {query[:100]}...") - raise ValueError( - f"Query execution timed out after {self.timeout} seconds in restricted mode. " - "Consider simplifying your query or increasing the timeout." - ) from e - except Exception as e: - logger.error(f"Error executing query: {e}") - raise - else: - return await self.sql_driver.execute_query( - f"/* crystaldba */ {query}", - params=params, - force_readonly=True, - ) - - @staticmethod - def sql_to_query(sql: Composable) -> str: - """Convert a SQL string to a query string.""" - return sql.as_string() - - @staticmethod - def param_sql_to_query(query: str, params: list[Any]) -> str: - """Convert a SQL string to a query string.""" - # Convert each parameter to a Literal only if it's not already a Composable - sql_params = [p if isinstance(p, Composable) else Literal(p) for p in params] - - return SafeSqlDriver.sql_to_query( - SQL(query).format(*sql_params) # type: ignore - ) - - @staticmethod - async def execute_param_query(sql_driver: SqlDriver, query: LiteralString, params: list[Any] | None = None) -> list[SqlDriver.RowResult] | None: - """Execute a query after validating it is safe""" - if params: - query_params = SafeSqlDriver.param_sql_to_query(query, params) - return await sql_driver.execute_query(query_params) # type: ignore - else: - return await sql_driver.execute_query(query) diff --git a/src/mcpg/_vendor/sql/sql_driver.py b/src/mcpg/_vendor/sql/sql_driver.py deleted file mode 100644 index cf96a13..0000000 --- a/src/mcpg/_vendor/sql/sql_driver.py +++ /dev/null @@ -1,276 +0,0 @@ -"""SQL driver adapter for PostgreSQL connections.""" - -import logging -import re -from dataclasses import dataclass -from typing import Any -from typing import Dict -from typing import List -from typing import Optional -from urllib.parse import urlparse -from urllib.parse import urlunparse - -from psycopg.rows import dict_row -from psycopg_pool import AsyncConnectionPool -from typing_extensions import LiteralString - -logger = logging.getLogger(__name__) - - -def obfuscate_password(text: str | None) -> str | None: - """ - Obfuscate password in any text containing connection information. - Works on connection URLs, error messages, and other strings. - """ - if text is None: - return None - - if not text: - return text - - # Try first as a proper URL - try: - parsed = urlparse(text) - if parsed.scheme and parsed.netloc and parsed.password: - # Replace password with asterisks in proper URL - netloc = parsed.netloc.replace(parsed.password, "****") - return urlunparse(parsed._replace(netloc=netloc)) - except Exception: - pass - - # Handle strings that contain connection strings but aren't proper URLs - # Match postgres://user:password@host:port/dbname pattern - url_pattern = re.compile(r"(postgres(?:ql)?:\/\/[^:]+:)([^@]+)(@[^\/\s]+)") - text = re.sub(url_pattern, r"\1****\3", text) - - # Match connection string parameters (password=xxx) - # This simpler pattern captures password without quotes - param_pattern = re.compile(r'(password=)([^\s&;"\']+)', re.IGNORECASE) - text = re.sub(param_pattern, r"\1****", text) - - # Match password in DSN format with single quotes - dsn_single_quote = re.compile(r"(password\s*=\s*')([^']+)(')", re.IGNORECASE) - text = re.sub(dsn_single_quote, r"\1****\3", text) - - # Match password in DSN format with double quotes - dsn_double_quote = re.compile(r'(password\s*=\s*")([^"]+)(")', re.IGNORECASE) - text = re.sub(dsn_double_quote, r"\1****\3", text) - - return text - - -class DbConnPool: - """Database connection manager using psycopg's connection pool.""" - - def __init__(self, connection_url: Optional[str] = None, min_size: int = 1, max_size: int = 5): - # MCPg local modification: min_size/max_size are configurable (ADR-0003). - # Defaults reproduce the original hardcoded behaviour. - self.connection_url = connection_url - self.min_size = min_size - self.max_size = max_size - self.pool: AsyncConnectionPool | None = None - self._is_valid = False - self._last_error = None - - async def pool_connect(self, connection_url: Optional[str] = None) -> AsyncConnectionPool: - """Initialize connection pool with retry logic.""" - # If we already have a valid pool, return it - if self.pool and self._is_valid: - return self.pool - - url = connection_url or self.connection_url - self.connection_url = url - if not url: - self._is_valid = False - self._last_error = "Database connection URL not provided" - raise ValueError(self._last_error) - - # Close any existing pool before creating a new one - await self.close() - - try: - # Configure connection pool with appropriate settings - self.pool = AsyncConnectionPool( - conninfo=url, - min_size=self.min_size, # MCPg local modification (ADR-0003) - max_size=self.max_size, # MCPg local modification (ADR-0003) - open=False, # Don't connect immediately, let's do it explicitly - ) - - # Open the pool explicitly - await self.pool.open() - - # Test the connection pool by executing a simple query - async with self.pool.connection() as conn: - async with conn.cursor() as cursor: - await cursor.execute("SELECT 1") - - self._is_valid = True - self._last_error = None - return self.pool - except Exception as e: - self._is_valid = False - self._last_error = str(e) - - # Clean up failed pool - await self.close() - - raise ValueError(f"Connection attempt failed: {obfuscate_password(str(e))}") from e - - async def close(self) -> None: - """Close the connection pool.""" - if self.pool: - try: - # Close the pool - await self.pool.close() - except Exception as e: - logger.warning(f"Error closing connection pool: {e}") - finally: - self.pool = None - self._is_valid = False - - @property - def is_valid(self) -> bool: - """Check if the connection pool is valid.""" - return self._is_valid - - @property - def last_error(self) -> Optional[str]: - """Get the last error message.""" - return self._last_error - - -class SqlDriver: - """Adapter class that wraps a PostgreSQL connection with the interface expected by DTA.""" - - @dataclass - class RowResult: - """Simple class to match the Griptape RowResult interface.""" - - cells: Dict[str, Any] - - def __init__( - self, - conn: Any = None, - engine_url: str | None = None, - ): - """ - Initialize with a PostgreSQL connection or pool. - - Args: - conn: PostgreSQL connection object or pool - engine_url: Connection URL string as an alternative to providing a connection - """ - if conn: - self.conn = conn - # Check if this is a connection pool - self.is_pool = isinstance(conn, DbConnPool) - elif engine_url: - # Don't connect here since we need async connection - self.engine_url = engine_url - self.conn = None - self.is_pool = False - else: - raise ValueError("Either conn or engine_url must be provided") - - def connect(self): - if self.conn is not None: - return self.conn - if self.engine_url: - self.conn = DbConnPool(self.engine_url) - self.is_pool = True - return self.conn - else: - raise ValueError("Connection not established. Either conn or engine_url must be provided") - - async def execute_query( - self, - query: LiteralString, - params: list[Any] | None = None, - force_readonly: bool = False, - ) -> Optional[List[RowResult]]: - """ - Execute a query and return results. - - Args: - query: SQL query to execute - params: Query parameters - force_readonly: Whether to enforce read-only mode - - Returns: - List of RowResult objects or None on error - """ - try: - if self.conn is None: - self.connect() - if self.conn is None: - raise ValueError("Connection not established") - - # Handle connection pool vs direct connection - if self.is_pool: - # For pools, get a connection from the pool - pool = await self.conn.pool_connect() - async with pool.connection() as connection: - return await self._execute_with_connection(connection, query, params, force_readonly=force_readonly) - else: - # Direct connection approach - return await self._execute_with_connection(self.conn, query, params, force_readonly=force_readonly) - except Exception as e: - # Mark pool as invalid if there was a connection issue - if self.conn and self.is_pool: - self.conn._is_valid = False # type: ignore - self.conn._last_error = str(e) # type: ignore - elif self.conn and not self.is_pool: - self.conn = None - - raise e - - async def _execute_with_connection(self, connection, query, params, force_readonly) -> Optional[List[RowResult]]: - """Execute query with the given connection.""" - transaction_started = False - try: - async with connection.cursor(row_factory=dict_row) as cursor: - # Start read-only transaction - if force_readonly: - await cursor.execute("BEGIN TRANSACTION READ ONLY") - transaction_started = True - - if params: - await cursor.execute(query, params) - else: - await cursor.execute(query) - - # For multiple statements, move to the last statement's results - while cursor.nextset(): - pass - - if cursor.description is None: # No results (like DDL statements) - if not force_readonly: - await cursor.execute("COMMIT") - elif transaction_started: - await cursor.execute("ROLLBACK") - transaction_started = False - return None - - # Get results from the last statement only - rows = await cursor.fetchall() - - # End the transaction appropriately - if not force_readonly: - await cursor.execute("COMMIT") - elif transaction_started: - await cursor.execute("ROLLBACK") - transaction_started = False - - return [SqlDriver.RowResult(cells=dict(row)) for row in rows] - - except Exception as e: - # Try to roll back the transaction if it's still active - if transaction_started: - try: - await connection.rollback() - except Exception as rollback_error: - logger.error(f"Error rolling back transaction: {rollback_error}") - - logger.error(f"Error executing query ({query}): {e}") - raise e diff --git a/src/mcpg/advisors.py b/src/mcpg/advisors.py index e29ac06..f91cddc 100644 --- a/src/mcpg/advisors.py +++ b/src/mcpg/advisors.py @@ -27,8 +27,8 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.query import QueryError, analyze_query_plan +from mcpg.sql import SqlDriver # Stable rule identifiers — agents may filter by these. RULE_MISSING_PRIMARY_KEY = "missing_primary_key" diff --git a/src/mcpg/aio.py b/src/mcpg/aio.py index 7fab7c8..fa4a1d9 100644 --- a/src/mcpg/aio.py +++ b/src/mcpg/aio.py @@ -43,7 +43,7 @@ from dataclasses import dataclass, field -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # PG 19 ships the AIO subsystem. The version-num probe is the boundary # guard — no extension to install. diff --git a/src/mcpg/audit.py b/src/mcpg/audit.py index 409576a..fe39f03 100644 --- a/src/mcpg/audit.py +++ b/src/mcpg/audit.py @@ -18,7 +18,7 @@ from dataclasses import dataclass, field from typing import Any -from mcpg._vendor.sql import SqlDriver, obfuscate_password +from mcpg.sql import SqlDriver, obfuscate_password # --- Tool invocation logger (Original Audit Logic) ------------------------ diff --git a/src/mcpg/audit_integrity.py b/src/mcpg/audit_integrity.py index fbb9345..a047904 100644 --- a/src/mcpg/audit_integrity.py +++ b/src/mcpg/audit_integrity.py @@ -13,7 +13,7 @@ from os import environ from typing import Any -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver AUDIT_SCHEMA = "mcpg_audit" AUDIT_TABLE = "events" diff --git a/src/mcpg/audit_nl2sql.py b/src/mcpg/audit_nl2sql.py index 67ad3da..032171b 100644 --- a/src/mcpg/audit_nl2sql.py +++ b/src/mcpg/audit_nl2sql.py @@ -47,8 +47,8 @@ from os import environ from typing import Any, Literal -from mcpg._vendor.sql import SqlDriver, obfuscate_password from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver, obfuscate_password logger = logging.getLogger(__name__) diff --git a/src/mcpg/audit_trail.py b/src/mcpg/audit_trail.py index bc0ba28..1282b11 100644 --- a/src/mcpg/audit_trail.py +++ b/src/mcpg/audit_trail.py @@ -26,7 +26,6 @@ from os import environ from typing import Any -from mcpg._vendor.sql import SqlDriver, obfuscate_password from mcpg.audit import is_secret_key from mcpg.audit_nl2sql import Backend from mcpg.audit_nl2sql import NL2SQLAuditError as _SharedAuditError @@ -35,6 +34,7 @@ from mcpg.audit_nl2sql import detect_backend as _detect_backend from mcpg.config import _parse_bool from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver, obfuscate_password _AUDIT_LOCK: asyncio.Lock | None = None diff --git a/src/mcpg/autovacuum.py b/src/mcpg/autovacuum.py index 10a8336..dae5243 100644 --- a/src/mcpg/autovacuum.py +++ b/src/mcpg/autovacuum.py @@ -46,7 +46,7 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # Cluster defaults that drive the per-table threshold formula. We don't # hardcode them — the query reads them from `current_setting` so a diff --git a/src/mcpg/composite.py b/src/mcpg/composite.py index a35563b..9f8bdf5 100644 --- a/src/mcpg/composite.py +++ b/src/mcpg/composite.py @@ -24,7 +24,6 @@ from dataclasses import dataclass, field from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( ColumnInfo, ConstraintInfo, @@ -37,6 +36,7 @@ ) from mcpg.liveops import list_active_queries from mcpg.query import analyze_query_plan, explain_query +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/config.py b/src/mcpg/config.py index 304915f..f511854 100644 --- a/src/mcpg/config.py +++ b/src/mcpg/config.py @@ -15,9 +15,9 @@ from os.path import isabs from urllib.parse import urlparse -from mcpg._vendor.sql import obfuscate_password from mcpg.nl2sql import AUTO_PICK_ORDER, VENDOR_ENV_VAR_HINT, VENDOR_KEY_ENV_VARS from mcpg.secrets import SecretsError, build_secrets_provider +from mcpg.sql import obfuscate_password # PG role names must be safe identifiers — we inline them into # ``SET ROLE ""`` so anything outside ``[A-Za-z_][A-Za-z0-9_]*`` diff --git a/src/mcpg/config_advisor.py b/src/mcpg/config_advisor.py index e3f13c3..6e58bd7 100644 --- a/src/mcpg/config_advisor.py +++ b/src/mcpg/config_advisor.py @@ -42,7 +42,7 @@ from dataclasses import dataclass, field -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # Status codes shared across the audit tools — same vocabulary as # mcpg.audit's MetricResult.status so an operator sees consistent diff --git a/src/mcpg/cron.py b/src/mcpg/cron.py index 7f275ae..15be877 100644 --- a/src/mcpg/cron.py +++ b/src/mcpg/cron.py @@ -12,8 +12,8 @@ import re from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver class CronError(Exception): diff --git a/src/mcpg/cursors.py b/src/mcpg/cursors.py index 685146a..5051291 100644 --- a/src/mcpg/cursors.py +++ b/src/mcpg/cursors.py @@ -28,7 +28,7 @@ import psycopg from psycopg.rows import dict_row -from mcpg._vendor.sql import SafeSqlDriver, SqlDriver +from mcpg.sql import SafeSqlDriver, SqlDriver logger = logging.getLogger(__name__) diff --git a/src/mcpg/data_movement.py b/src/mcpg/data_movement.py index f0aa9aa..ced6fb8 100644 --- a/src/mcpg/data_movement.py +++ b/src/mcpg/data_movement.py @@ -28,10 +28,10 @@ from typing import Any from urllib.parse import unquote, urlparse -from mcpg._vendor.sql import SqlDriver from mcpg.database import Database from mcpg.query import QueryError, run_select from mcpg.shell import ShellError, SubprocessLimits, run_pg_binary +from mcpg.sql import SqlDriver # Same identifier allowlist as mcpg.textsearch / mcpg.prisma / mcpg.vector_tuning — # refuse names that need delimited-identifier quoting, accept plain ones. diff --git a/src/mcpg/database.py b/src/mcpg/database.py index a67370b..4e487cb 100644 --- a/src/mcpg/database.py +++ b/src/mcpg/database.py @@ -12,10 +12,10 @@ from types import TracebackType from typing import Any -from mcpg._vendor.sql import DbConnPool, SqlDriver, obfuscate_password from mcpg.config import Settings from mcpg.multidb import PRIMARY_DATABASE_ID, make_read_only_driver from mcpg.replicas import ReplicaPool, RoutedSqlDriver, _make_driver_for_pool +from mcpg.sql import DbConnPool, SqlDriver, obfuscate_password logger = logging.getLogger(__name__) diff --git a/src/mcpg/diagrams.py b/src/mcpg/diagrams.py index 1563328..e935e58 100644 --- a/src/mcpg/diagrams.py +++ b/src/mcpg/diagrams.py @@ -10,13 +10,13 @@ import re -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( describe_table, list_constraints, list_foreign_keys, list_tables, ) +from mcpg.sql import SqlDriver _PRIMARY_KEY_COLUMNS = re.compile(r"PRIMARY KEY \(([^)]+)\)", re.IGNORECASE) _NON_IDENT = re.compile(r"[^A-Za-z0-9]+") diff --git a/src/mcpg/diesel.py b/src/mcpg/diesel.py index 41938d9..079ed57 100644 --- a/src/mcpg/diesel.py +++ b/src/mcpg/diesel.py @@ -21,7 +21,6 @@ import re -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( ColumnInfo, describe_table, @@ -30,6 +29,7 @@ list_foreign_keys, list_tables, ) +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/drizzle.py b/src/mcpg/drizzle.py index 0a3a9be..9b5999c 100644 --- a/src/mcpg/drizzle.py +++ b/src/mcpg/drizzle.py @@ -17,7 +17,6 @@ import re from collections.abc import Iterable -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( ColumnInfo, ForeignKeyInfo, @@ -29,6 +28,7 @@ list_indexes, list_tables, ) +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/ecto.py b/src/mcpg/ecto.py index f981e04..013f051 100644 --- a/src/mcpg/ecto.py +++ b/src/mcpg/ecto.py @@ -20,7 +20,6 @@ import re -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( ColumnInfo, ForeignKeyInfo, @@ -30,6 +29,7 @@ list_foreign_keys, list_tables, ) +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/ent.py b/src/mcpg/ent.py index 872654f..8459a25 100644 --- a/src/mcpg/ent.py +++ b/src/mcpg/ent.py @@ -17,7 +17,6 @@ import re -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( ColumnInfo, ForeignKeyInfo, @@ -26,6 +25,7 @@ list_foreign_keys, list_tables, ) +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/extensions.py b/src/mcpg/extensions.py index 62563cf..5f864bf 100644 --- a/src/mcpg/extensions.py +++ b/src/mcpg/extensions.py @@ -10,7 +10,7 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # Extensions MCPg will enable on request: well-known, widely-used extensions. # CREATE EXTENSION takes an identifier, so this allowlist guards against diff --git a/src/mcpg/graph_projection.py b/src/mcpg/graph_projection.py index 164908a..0204fac 100644 --- a/src/mcpg/graph_projection.py +++ b/src/mcpg/graph_projection.py @@ -39,8 +39,8 @@ import re from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ColumnInfo, describe_table, list_foreign_keys +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"\A[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/headline_curator.py b/src/mcpg/headline_curator.py index 910b19c..464837b 100644 --- a/src/mcpg/headline_curator.py +++ b/src/mcpg/headline_curator.py @@ -30,8 +30,8 @@ from dataclasses import dataclass, field -from mcpg._vendor.sql import SqlDriver from mcpg.about import CAPABILITIES, classify_tool +from mcpg.sql import SqlDriver class HeadlineCuratorError(Exception): diff --git a/src/mcpg/health.py b/src/mcpg/health.py index f69f5bc..633e02c 100644 --- a/src/mcpg/health.py +++ b/src/mcpg/health.py @@ -9,8 +9,8 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver # Classification thresholds. _CONNECTION_WARN_RATIO = 0.8 # warn above 80% of max_connections diff --git a/src/mcpg/indexing.py b/src/mcpg/indexing.py index 9b61465..7a2507d 100644 --- a/src/mcpg/indexing.py +++ b/src/mcpg/indexing.py @@ -15,7 +15,7 @@ from dataclasses import dataclass, field -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # Tables smaller than this are ignored — sequential scans of them are cheap. DEFAULT_MIN_LIVE_TUPLES = 10_000 diff --git a/src/mcpg/introspection.py b/src/mcpg/introspection.py index e95e3f0..3516419 100644 --- a/src/mcpg/introspection.py +++ b/src/mcpg/introspection.py @@ -10,7 +10,7 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # Schemas that belong to PostgreSQL itself rather than the user. _SYSTEM_SCHEMAS = frozenset({"pg_catalog", "information_schema", "pg_toast"}) diff --git a/src/mcpg/io_stats.py b/src/mcpg/io_stats.py index abb3292..ab74890 100644 --- a/src/mcpg/io_stats.py +++ b/src/mcpg/io_stats.py @@ -14,7 +14,7 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver @dataclass(frozen=True) diff --git a/src/mcpg/jooq.py b/src/mcpg/jooq.py index 510fef8..512418d 100644 --- a/src/mcpg/jooq.py +++ b/src/mcpg/jooq.py @@ -34,12 +34,12 @@ import re from xml.sax.saxutils import escape -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( ColumnInfo, describe_table, list_tables, ) +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/liveops.py b/src/mcpg/liveops.py index cac21ca..8ad2208 100644 --- a/src/mcpg/liveops.py +++ b/src/mcpg/liveops.py @@ -8,7 +8,7 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver @dataclass(frozen=True) diff --git a/src/mcpg/locks.py b/src/mcpg/locks.py index 8943ebb..22679d0 100644 --- a/src/mcpg/locks.py +++ b/src/mcpg/locks.py @@ -17,7 +17,7 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver class LocksError(Exception): diff --git a/src/mcpg/migration_history.py b/src/mcpg/migration_history.py index 4b45cca..60f6c8a 100644 --- a/src/mcpg/migration_history.py +++ b/src/mcpg/migration_history.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver @dataclass(frozen=True) diff --git a/src/mcpg/migration_ingestion.py b/src/mcpg/migration_ingestion.py index c4c1495..524bb7b 100644 --- a/src/mcpg/migration_ingestion.py +++ b/src/mcpg/migration_ingestion.py @@ -42,7 +42,7 @@ from dataclasses import dataclass, field from pathlib import Path -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver _FRAMEWORKS = frozenset({"alembic", "flyway", "liquibase"}) diff --git a/src/mcpg/migrations.py b/src/mcpg/migrations.py index b2cf7a5..31fb05f 100644 --- a/src/mcpg/migrations.py +++ b/src/mcpg/migrations.py @@ -25,9 +25,9 @@ from datetime import UTC, datetime, timedelta from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import describe_table, list_constraints, list_indexes, list_tables from mcpg.schema_diff import SchemaDiff, compare_schemas +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") _SHADOW_PREFIX = "mcpg_shadow_" diff --git a/src/mcpg/multidb.py b/src/mcpg/multidb.py index c950bff..543cbc7 100644 --- a/src/mcpg/multidb.py +++ b/src/mcpg/multidb.py @@ -27,7 +27,7 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import DbConnPool, SqlDriver +from mcpg.sql import DbConnPool, SqlDriver # The implicit id of ``MCPG_DATABASE_URL``. Reserved (see config.py) so a # secondary can never shadow it. diff --git a/src/mcpg/naming.py b/src/mcpg/naming.py index 94128cf..9d30670 100644 --- a/src/mcpg/naming.py +++ b/src/mcpg/naming.py @@ -28,7 +28,7 @@ from collections.abc import Iterable from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver DEFAULT_INDEX_PREFIXES: tuple[str, ...] = ("idx_", "ix_", "pk_", "uq_", "fk_", "gin_", "gist_", "brin_", "hnsw_") diff --git a/src/mcpg/nl2sql.py b/src/mcpg/nl2sql.py index 98fd785..bbe8f79 100644 --- a/src/mcpg/nl2sql.py +++ b/src/mcpg/nl2sql.py @@ -38,9 +38,9 @@ import httpx -from mcpg._vendor.sql import SqlDriver, obfuscate_password from mcpg.introspection import describe_table, list_foreign_keys, list_tables from mcpg.query import DEFAULT_MAX_ROWS, QueryError, explain_query, run_select +from mcpg.sql import SqlDriver, obfuscate_password if TYPE_CHECKING: from mcpg.config import Settings diff --git a/src/mcpg/otel_tracing.py b/src/mcpg/otel_tracing.py index c689cd7..eb0c2f2 100644 --- a/src/mcpg/otel_tracing.py +++ b/src/mcpg/otel_tracing.py @@ -36,7 +36,7 @@ from collections.abc import Iterator from typing import TYPE_CHECKING, Any -from mcpg._vendor.sql import obfuscate_password +from mcpg.sql import obfuscate_password if TYPE_CHECKING: from mcpg.config import Settings diff --git a/src/mcpg/partman.py b/src/mcpg/partman.py index e9751db..74cd519 100644 --- a/src/mcpg/partman.py +++ b/src/mcpg/partman.py @@ -13,8 +13,8 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver # pg_partman accepts these partition-type strings since version 5.x — # 'range' and 'list' for declarative partitioning, 'native' as a legacy diff --git a/src/mcpg/pg19_ddl.py b/src/mcpg/pg19_ddl.py index 7751766..0e19e19 100644 --- a/src/mcpg/pg19_ddl.py +++ b/src/mcpg/pg19_ddl.py @@ -38,8 +38,8 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver from mcpg.database import Database +from mcpg.sql import SqlDriver # PG 19 ships the new pg_get_*def() functions. The version-num boundary. _MIN_PG19_DDL_VERSION = 190000 diff --git a/src/mcpg/pg19_partitions.py b/src/mcpg/pg19_partitions.py index fe8917e..6047f3c 100644 --- a/src/mcpg/pg19_partitions.py +++ b/src/mcpg/pg19_partitions.py @@ -58,8 +58,8 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver from mcpg.database import Database +from mcpg.sql import SqlDriver # PG 19 ships both forms. The version-num boundary. _MIN_PG19_PARTITIONS_VERSION = 190000 diff --git a/src/mcpg/pg19_runtime.py b/src/mcpg/pg19_runtime.py index e9d7838..4a6bf60 100644 --- a/src/mcpg/pg19_runtime.py +++ b/src/mcpg/pg19_runtime.py @@ -41,8 +41,8 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver from mcpg.database import Database +from mcpg.sql import SqlDriver # PG 19 ships both toggles. The version-num probe is the boundary — # no extension to install. diff --git a/src/mcpg/pg19_skip_scan.py b/src/mcpg/pg19_skip_scan.py index d8184d0..934c79a 100644 --- a/src/mcpg/pg19_skip_scan.py +++ b/src/mcpg/pg19_skip_scan.py @@ -46,7 +46,7 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # PG 19 ships skip-scan as the planner default. The version-num boundary. _MIN_PG19_SKIP_SCAN_VERSION = 190000 diff --git a/src/mcpg/pg19_stats.py b/src/mcpg/pg19_stats.py index ffa9c72..c321152 100644 --- a/src/mcpg/pg19_stats.py +++ b/src/mcpg/pg19_stats.py @@ -36,7 +36,7 @@ from dataclasses import dataclass, field -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # Both views landed in PG 19. The version-num probe is the boundary — # no extension to install. diff --git a/src/mcpg/pg_prewarm.py b/src/mcpg/pg_prewarm.py index b9ba489..0802592 100644 --- a/src/mcpg/pg_prewarm.py +++ b/src/mcpg/pg_prewarm.py @@ -32,8 +32,8 @@ from dataclasses import dataclass, field -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver # pg_prewarm modes (matches the upstream signature: pg_prewarm(regclass, # mode text, fork text, first_block, last_block)). We don't expose ``fork`` diff --git a/src/mcpg/pg_search.py b/src/mcpg/pg_search.py index 7a61e1f..fd7a96f 100644 --- a/src/mcpg/pg_search.py +++ b/src/mcpg/pg_search.py @@ -66,9 +66,9 @@ from dataclasses import dataclass, field from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.database import Database from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver # Plain unquoted PostgreSQL identifier — same rule as turboquant / # vector_tuning. Anything that would require delimited quoting is diff --git a/src/mcpg/pgq.py b/src/mcpg/pgq.py index 8fef16e..dc5834f 100644 --- a/src/mcpg/pgq.py +++ b/src/mcpg/pgq.py @@ -52,7 +52,7 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # SQL/PGQ landed in the PG 19 series. We use the version-num probe rather # than an extension allowlist because SQL/PGQ is built into the server, diff --git a/src/mcpg/pitr.py b/src/mcpg/pitr.py index 479468c..ce06d6d 100644 --- a/src/mcpg/pitr.py +++ b/src/mcpg/pitr.py @@ -27,7 +27,7 @@ from dataclasses import dataclass, field -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver from mcpg.wal_archive import get_wal_archive_status # wal_level ordering for the ">= replica" gate. diff --git a/src/mcpg/prisma.py b/src/mcpg/prisma.py index d8a4a49..e898f68 100644 --- a/src/mcpg/prisma.py +++ b/src/mcpg/prisma.py @@ -29,7 +29,6 @@ import re from collections.abc import Iterable -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( ColumnInfo, EnumInfo, @@ -42,6 +41,7 @@ list_indexes, list_tables, ) +from mcpg.sql import SqlDriver # Generic PG types → Prisma scalar types. Types with parameters (e.g. # ``character varying(255)``, ``numeric(10,2)``, ``vector(384)``) are diff --git a/src/mcpg/query.py b/src/mcpg/query.py index 880fe33..d0c9458 100644 --- a/src/mcpg/query.py +++ b/src/mcpg/query.py @@ -15,7 +15,7 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import SafeSqlDriver, SqlDriver +from mcpg.sql import SafeSqlDriver, SqlDriver # Default per-query execution timeout, in seconds. DEFAULT_TIMEOUT_SECONDS = 30.0 diff --git a/src/mcpg/rag_efficiency.py b/src/mcpg/rag_efficiency.py index 4abb187..989554c 100644 --- a/src/mcpg/rag_efficiency.py +++ b/src/mcpg/rag_efficiency.py @@ -49,8 +49,8 @@ from dataclasses import dataclass, field from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver # Reused from vector_tuning's identifier convention — refuse anything # that would require delimited quoting. diff --git a/src/mcpg/rag_telemetry.py b/src/mcpg/rag_telemetry.py index 8661c45..b237cea 100644 --- a/src/mcpg/rag_telemetry.py +++ b/src/mcpg/rag_telemetry.py @@ -32,13 +32,13 @@ from os import environ from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.audit_nl2sql import Backend from mcpg.audit_nl2sql import NL2SQLAuditError as _SharedAuditError from mcpg.audit_nl2sql import _check_identifier as _shared_check_identifier from mcpg.audit_nl2sql import _check_interval as _shared_check_interval from mcpg.audit_nl2sql import detect_backend as _detect_backend from mcpg.database import Database +from mcpg.sql import SqlDriver _SCHEMA_NAME = "mcpg_rag" _TABLE_NAME = "rerank_events" diff --git a/src/mcpg/redis_fdw.py b/src/mcpg/redis_fdw.py index 69d5d45..e643260 100644 --- a/src/mcpg/redis_fdw.py +++ b/src/mcpg/redis_fdw.py @@ -37,10 +37,10 @@ from collections.abc import Mapping from dataclasses import dataclass, field -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed from mcpg.introspection import _parse_options from mcpg.secrets import SecretsProvider +from mcpg.sql import SqlDriver # redis_fdw identifies itself in pg_foreign_data_wrapper as "redis_fdw" — # we use that as the discriminator on every catalog filter below. diff --git a/src/mcpg/repack.py b/src/mcpg/repack.py index abc64b0..a78f4df 100644 --- a/src/mcpg/repack.py +++ b/src/mcpg/repack.py @@ -39,8 +39,8 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver from mcpg.database import Database +from mcpg.sql import SqlDriver # PG 19 ships REPACK; older versions don't recognise the keyword. _MIN_REPACK_VERSION = 190000 diff --git a/src/mcpg/replicas.py b/src/mcpg/replicas.py index 5f5546d..bf075ef 100644 --- a/src/mcpg/replicas.py +++ b/src/mcpg/replicas.py @@ -37,7 +37,7 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import DbConnPool, SqlDriver, obfuscate_password +from mcpg.sql import DbConnPool, SqlDriver, obfuscate_password from mcpg.tenancy import TenantSqlDriver logger = logging.getLogger(__name__) diff --git a/src/mcpg/rls.py b/src/mcpg/rls.py index 911c3bd..2d87604 100644 --- a/src/mcpg/rls.py +++ b/src/mcpg/rls.py @@ -24,7 +24,7 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"\A[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/schema_diff.py b/src/mcpg/schema_diff.py index 61f90aa..914d071 100644 --- a/src/mcpg/schema_diff.py +++ b/src/mcpg/schema_diff.py @@ -12,7 +12,6 @@ from collections.abc import Callable from dataclasses import dataclass, fields -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( ColumnInfo, ConstraintInfo, @@ -25,6 +24,7 @@ list_indexes, list_tables, ) +from mcpg.sql import SqlDriver @dataclass(frozen=True) diff --git a/src/mcpg/schema_docs.py b/src/mcpg/schema_docs.py index a9af240..53775d0 100644 --- a/src/mcpg/schema_docs.py +++ b/src/mcpg/schema_docs.py @@ -10,7 +10,6 @@ import re from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( describe_table, list_constraints, @@ -22,6 +21,7 @@ list_tables, list_views, ) +from mcpg.sql import SqlDriver _PRIMARY_KEY_COLUMNS = re.compile(r"PRIMARY KEY \(([^)]+)\)", re.IGNORECASE) diff --git a/src/mcpg/session_advisor.py b/src/mcpg/session_advisor.py index 003c549..14cfb4f 100644 --- a/src/mcpg/session_advisor.py +++ b/src/mcpg/session_advisor.py @@ -36,7 +36,7 @@ from dataclasses import dataclass, field -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # Tools that return small, well-defined slices of the catalogue — calling # them many times in a session is the canonical "would have been one diff --git a/src/mcpg/sqlalchemy_export.py b/src/mcpg/sqlalchemy_export.py index fbc10e7..e62c8bd 100644 --- a/src/mcpg/sqlalchemy_export.py +++ b/src/mcpg/sqlalchemy_export.py @@ -16,7 +16,6 @@ import re -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( ColumnInfo, ForeignKeyInfo, @@ -26,6 +25,7 @@ list_foreign_keys, list_tables, ) +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") _SERIAL_DEFAULT_RE = re.compile(r"nextval\(['\"]([^'\"]+)['\"]") diff --git a/src/mcpg/sqlc.py b/src/mcpg/sqlc.py index 5eb1c0b..4e4d548 100644 --- a/src/mcpg/sqlc.py +++ b/src/mcpg/sqlc.py @@ -22,7 +22,6 @@ import re -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ( ColumnInfo, describe_table, @@ -32,6 +31,7 @@ list_indexes, list_tables, ) +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/tenancy.py b/src/mcpg/tenancy.py index 88531fd..58e31bd 100644 --- a/src/mcpg/tenancy.py +++ b/src/mcpg/tenancy.py @@ -32,7 +32,7 @@ from psycopg.rows import dict_row -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver logger = logging.getLogger(__name__) diff --git a/src/mcpg/test_data.py b/src/mcpg/test_data.py index 402f6b9..2f22784 100644 --- a/src/mcpg/test_data.py +++ b/src/mcpg/test_data.py @@ -32,8 +32,8 @@ from dataclasses import dataclass from datetime import UTC, datetime, timedelta -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ColumnInfo, describe_table +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"\A[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/test_row_factory.py b/src/mcpg/test_row_factory.py index ff3bf40..42a3fe3 100644 --- a/src/mcpg/test_row_factory.py +++ b/src/mcpg/test_row_factory.py @@ -58,8 +58,8 @@ from dataclasses import dataclass from datetime import UTC, datetime, timedelta -from mcpg._vendor.sql import SqlDriver from mcpg.introspection import ColumnInfo, describe_table +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"\A[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/textsearch.py b/src/mcpg/textsearch.py index ab0e000..9c3df7b 100644 --- a/src/mcpg/textsearch.py +++ b/src/mcpg/textsearch.py @@ -19,8 +19,8 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/timescaledb.py b/src/mcpg/timescaledb.py index 0ed3635..b76f70c 100644 --- a/src/mcpg/timescaledb.py +++ b/src/mcpg/timescaledb.py @@ -19,8 +19,8 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") # A whitelist of intervals we'll inline into SQL. Validated against this diff --git a/src/mcpg/tools.py b/src/mcpg/tools.py index 0f968a6..46276c1 100644 --- a/src/mcpg/tools.py +++ b/src/mcpg/tools.py @@ -93,10 +93,10 @@ workload, write, ) -from mcpg._vendor.sql import SqlDriver from mcpg.config import Settings from mcpg.context import AppContext from mcpg.policy import Capability, is_permitted +from mcpg.sql import SqlDriver # The MCP request context FastMCP injects into every tool. _Ctx = Context[ServerSession, AppContext, Any] diff --git a/src/mcpg/turboquant.py b/src/mcpg/turboquant.py index 3e5714c..5143cfa 100644 --- a/src/mcpg/turboquant.py +++ b/src/mcpg/turboquant.py @@ -58,9 +58,9 @@ from dataclasses import dataclass, field from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.database import Database from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver # Plain unquoted PostgreSQL identifier — matches the rule used by # vector_tuning. Anything that would require delimited quoting at the diff --git a/src/mcpg/vector_ops.py b/src/mcpg/vector_ops.py index 28dc61d..85c6086 100644 --- a/src/mcpg/vector_ops.py +++ b/src/mcpg/vector_ops.py @@ -24,8 +24,8 @@ from typing import Any from mcpg import introspection -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver _IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") diff --git a/src/mcpg/vector_tuner_advanced.py b/src/mcpg/vector_tuner_advanced.py index f83f4b0..b47feb0 100644 --- a/src/mcpg/vector_tuner_advanced.py +++ b/src/mcpg/vector_tuner_advanced.py @@ -23,8 +23,8 @@ from dataclasses import dataclass, field from typing import Any -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver from mcpg.vector_tuning import VectorTuningError, _indexes_on_column, _quoted _DISTANCE_OPERATORS = {"l2": "<->", "cosine": "<=>", "inner_product": "<#>"} diff --git a/src/mcpg/vector_tuning.py b/src/mcpg/vector_tuning.py index 8a03b41..927d4ab 100644 --- a/src/mcpg/vector_tuning.py +++ b/src/mcpg/vector_tuning.py @@ -19,9 +19,9 @@ import re from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed from mcpg.introspection import describe_table +from mcpg.sql import SqlDriver # pgvector supports these index access methods today; the allowlist # guards against arbitrary identifier injection in CREATE INDEX text. diff --git a/src/mcpg/wait_for_lsn.py b/src/mcpg/wait_for_lsn.py index 7305a3d..5d56348 100644 --- a/src/mcpg/wait_for_lsn.py +++ b/src/mcpg/wait_for_lsn.py @@ -50,7 +50,7 @@ import re from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver # PG 19 ships WAIT FOR LSN. The version-num boundary. _MIN_PG19_WAIT_VERSION = 190000 diff --git a/src/mcpg/wal_archive.py b/src/mcpg/wal_archive.py index e2969f3..5fa0763 100644 --- a/src/mcpg/wal_archive.py +++ b/src/mcpg/wal_archive.py @@ -34,7 +34,7 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver @dataclass(frozen=True) diff --git a/src/mcpg/walinspect.py b/src/mcpg/walinspect.py index cb56de4..ba71928 100644 --- a/src/mcpg/walinspect.py +++ b/src/mcpg/walinspect.py @@ -4,7 +4,7 @@ from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver @dataclass(frozen=True) diff --git a/src/mcpg/warehousepg.py b/src/mcpg/warehousepg.py index afcfab7..261eefb 100644 --- a/src/mcpg/warehousepg.py +++ b/src/mcpg/warehousepg.py @@ -66,7 +66,7 @@ from dataclasses import dataclass from typing import Any -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver @dataclass(frozen=True) diff --git a/src/mcpg/workload.py b/src/mcpg/workload.py index c02ecb7..9db1f0a 100644 --- a/src/mcpg/workload.py +++ b/src/mcpg/workload.py @@ -16,8 +16,8 @@ import re from dataclasses import dataclass -from mcpg._vendor.sql import SqlDriver from mcpg.extensions import extension_installed +from mcpg.sql import SqlDriver # Default number of slow queries to return. DEFAULT_LIMIT = 10 diff --git a/src/mcpg/write.py b/src/mcpg/write.py index 2bec245..dd4424f 100644 --- a/src/mcpg/write.py +++ b/src/mcpg/write.py @@ -21,8 +21,8 @@ import pglast -from mcpg._vendor.sql import SqlDriver from mcpg.audit_trail import SchemaDiffSnapshot, capture_columns, record_audit +from mcpg.sql import SqlDriver # pglast statement node names accepted by run_write. _DML_STATEMENTS = frozenset({"InsertStmt", "UpdateStmt", "DeleteStmt"}) diff --git a/tests/unit/_fakes.py b/tests/unit/_fakes.py index b7820fa..4eccb66 100644 --- a/tests/unit/_fakes.py +++ b/tests/unit/_fakes.py @@ -4,7 +4,7 @@ from typing import Any -from mcpg._vendor.sql import SqlDriver +from mcpg.sql import SqlDriver class FakePool: diff --git a/tests/unit/test_audit_integrity.py b/tests/unit/test_audit_integrity.py index 81445b0..f0413e3 100644 --- a/tests/unit/test_audit_integrity.py +++ b/tests/unit/test_audit_integrity.py @@ -454,7 +454,7 @@ def __init__(self) -> None: self.calls: list[tuple[str, list[Any] | None, bool]] = [] async def execute_query(self, query: str, params: list[Any] | None = None, force_readonly: bool = False) -> Any: - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver self.calls.append((query, params, force_readonly)) if "FROM pg_class c" in query: diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index 8fa5f75..9109008 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -3,9 +3,9 @@ import pytest from _fakes import FakePool -from mcpg._vendor.sql import DbConnPool, SqlDriver from mcpg.config import load_settings from mcpg.database import Database, DatabaseError +from mcpg.sql import DbConnPool, SqlDriver _SETTINGS = load_settings({"MCPG_DATABASE_URL": "postgresql://u:p@localhost/db"}) diff --git a/tests/unit/test_graph_projection.py b/tests/unit/test_graph_projection.py index 5647e1a..5aedf7b 100644 --- a/tests/unit/test_graph_projection.py +++ b/tests/unit/test_graph_projection.py @@ -7,7 +7,6 @@ import pytest from _fakes import FakeRoutingDriver -from mcpg._vendor.sql import SqlDriver from mcpg.graph_projection import ( HARD_ROW_CAP, EdgeType, @@ -15,6 +14,7 @@ NodeLabel, generate_graph_projection, ) +from mcpg.sql import SqlDriver class _AgeAwareDriver(FakeRoutingDriver): diff --git a/tests/unit/test_introspection.py b/tests/unit/test_introspection.py index 8b83e4e..06ce963 100644 --- a/tests/unit/test_introspection.py +++ b/tests/unit/test_introspection.py @@ -471,7 +471,7 @@ async def execute_query( self.calls.append((query, params)) if "pg_get_acl" in query: raise RuntimeError("function pg_get_acl(regclass, oid, integer) does not exist") - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver return [ SqlDriver.RowResult( diff --git a/tests/unit/test_migration_history.py b/tests/unit/test_migration_history.py index d47a497..f2296c4 100644 --- a/tests/unit/test_migration_history.py +++ b/tests/unit/test_migration_history.py @@ -8,7 +8,6 @@ from _fakes import FakeDatabase, FakeRoutingDriver from mcp.shared.memory import create_connected_server_and_client_session -from mcpg._vendor.sql import SqlDriver from mcpg.config import load_settings from mcpg.migration_history import ( AlembicMigration, @@ -23,6 +22,7 @@ read_migration_history, ) from mcpg.server import create_server +from mcpg.sql import SqlDriver _SETTINGS = load_settings({"MCPG_DATABASE_URL": "postgresql://u:p@localhost/db"}) diff --git a/tests/unit/test_multidb.py b/tests/unit/test_multidb.py index 0e99385..7386aab 100644 --- a/tests/unit/test_multidb.py +++ b/tests/unit/test_multidb.py @@ -17,7 +17,6 @@ import pytest from _fakes import FakePool -from mcpg._vendor.sql import SqlDriver from mcpg.config import load_settings from mcpg.database import Database, DatabaseError from mcpg.multidb import ( @@ -27,6 +26,7 @@ make_read_only_driver, ) from mcpg.replicas import TimeoutSqlDriver +from mcpg.sql import SqlDriver _PRIMARY = "postgresql://u:p@localhost/db" diff --git a/tests/unit/test_pg19_ddl.py b/tests/unit/test_pg19_ddl.py index 8993597..174ade9 100644 --- a/tests/unit/test_pg19_ddl.py +++ b/tests/unit/test_pg19_ddl.py @@ -81,7 +81,7 @@ def __init__(self, *, convalidated: bool | None, alter_fails: bool = False) -> N self.calls: list[tuple[str, object, bool]] = [] async def execute_query(self, query, params=None, force_readonly=False): # type: ignore[no-untyped-def] - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver self.calls.append((query, params, force_readonly)) if "pg_constraint" in query: diff --git a/tests/unit/test_pg19_skip_scan.py b/tests/unit/test_pg19_skip_scan.py index b409bbb..62a76d2 100644 --- a/tests/unit/test_pg19_skip_scan.py +++ b/tests/unit/test_pg19_skip_scan.py @@ -99,7 +99,7 @@ def __init__(self) -> None: self.calls: list[str] = [] async def execute_query(self, query, params=None, force_readonly=False): # type: ignore[no-untyped-def] - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver self.calls.append(query) if "current_setting" in query: diff --git a/tests/unit/test_sql_kernel_differential.py b/tests/unit/test_sql_kernel_differential.py deleted file mode 100644 index 66b1b8f..0000000 --- a/tests/unit/test_sql_kernel_differential.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Differential parity harness — first-party validator vs the vendored one. - -The security-review gate for the de-vendor effort (roadmap 18.1): feed a -broad corpus of safe / unsafe / malformed SQL through **both** -``mcpg._vendor.sql.SafeSqlDriver`` and the first-party -``mcpg.sql.SafeSqlDriver`` and assert an **identical accept/reject verdict -on every input**. Zero divergence is the bar — this proves the re-author -didn't move the security boundary. - -Temporary: deleted in PR 2 together with ``_vendor/``. -""" - -from __future__ import annotations - -from unittest.mock import Mock - -import pytest - -from mcpg._vendor.sql import SafeSqlDriver as VendoredSafeSqlDriver -from mcpg.sql import SafeSqlDriver as FirstPartySafeSqlDriver - -# Corpus: a mix of accepted read-only statements, rejected write/DDL/unsafe -# statements, and malformed / adversarial inputs. The point is not whether a -# given item is accepted, but that BOTH validators agree on it. -_CORPUS: list[str] = [ - # --- expected-safe read-only statements --- - "SELECT 1", - "SELECT * FROM users", - "SELECT id, name FROM public.users WHERE id = 5", - "SELECT count(*), max(created_at) FROM orders GROUP BY user_id", - "SELECT * FROM a JOIN b ON a.id = b.a_id WHERE a.name LIKE 'foo%'", - "WITH t AS (SELECT id FROM users) SELECT * FROM t", - "SELECT * FROM users ORDER BY id LIMIT 10 OFFSET 5", - "SELECT lower(name), upper(email) FROM users", - "SELECT * FROM users u WHERE u.id IN (SELECT user_id FROM orders)", - "EXPLAIN SELECT * FROM users", - "EXPLAIN (FORMAT JSON) SELECT * FROM users", - "SHOW search_path", - "VACUUM users", - "ANALYZE users", - "SELECT jsonb_agg(row_to_json(u)) FROM users u", - "SELECT array_agg(id) FROM users", - "SELECT coalesce(name, 'n/a') FROM users", - # --- expected-unsafe: writes / DDL / DCL --- - "INSERT INTO users (name) VALUES ('x')", - "UPDATE users SET name = 'x' WHERE id = 1", - "DELETE FROM users WHERE id = 1", - "DROP TABLE users", - "CREATE TABLE t (id int)", - "ALTER TABLE users ADD COLUMN x int", - "TRUNCATE users", - "GRANT SELECT ON users TO public", - "CREATE EXTENSION pg_stat_statements", - "CREATE EXTENSION definitely_not_allowed_ext", - "COPY users TO '/tmp/x.csv'", - "COPY users FROM '/tmp/x.csv'", - "DO $$ BEGIN PERFORM 1; END $$", - "SET search_path TO public", - # --- expected-unsafe: read-only escapes --- - "SELECT * FROM users FOR UPDATE", - "EXPLAIN ANALYZE SELECT * FROM users", - "SELECT pg_sleep(10)", - "SELECT * FROM pg_read_file('/etc/passwd')", - "SELECT lo_import('/etc/passwd')", - "SELECT 1; DROP TABLE users", - "SELECT * FROM users; SELECT * FROM orders", - # --- malformed / adversarial / edge --- - "", - " ", - "-- just a comment", - "SELECT", - "SELECT * FROM", - "SELECT * FROM users WHERE name = 'unterminated", - "SEL ECT bogus", - "SELECT /* nested */ 1", - "SELECT 'foo'; -- SELECT 2", - "SELECT * FROM üsers", # unicode identifier - "SELECT " + "1 + " * 50 + "1", # deeply nested expr -] - - -def _verdict(driver_cls: type, sql: str) -> tuple[bool, str]: - """Return (accepted, error-class-name) for one validator on one input. - - ``accepted`` is ``True`` when ``_validate`` does not raise. - """ - driver = driver_cls(Mock()) - try: - driver._validate(sql) - return (True, "") - except Exception as exc: # we compare verdicts, not raise - return (False, type(exc).__name__) - - -@pytest.mark.parametrize("sql", _CORPUS, ids=lambda s: s[:40] or "") -def test_first_party_validator_matches_vendored_verdict(sql: str) -> None: - vendored_ok, vendored_err = _verdict(VendoredSafeSqlDriver, sql) - first_party_ok, first_party_err = _verdict(FirstPartySafeSqlDriver, sql) - - assert first_party_ok == vendored_ok, ( - f"verdict divergence on {sql!r}: " - f"vendored={'accept' if vendored_ok else 'reject'} " - f"first-party={'accept' if first_party_ok else 'reject'}" - ) - # On rejection, the raised exception class should match too (both raise - # ValueError from the validator). - if not vendored_ok: - assert first_party_err == vendored_err, ( - f"error-type divergence on {sql!r}: vendored={vendored_err} first-party={first_party_err}" - ) - - -def test_corpus_exercises_both_verdicts() -> None: - """Sanity: the corpus contains both accepted and rejected inputs, so the - parity test isn't vacuously passing on an all-reject (or all-accept) set.""" - verdicts = {_verdict(FirstPartySafeSqlDriver, sql)[0] for sql in _CORPUS} - assert verdicts == {True, False} diff --git a/tests/unit/test_tools.py b/tests/unit/test_tools.py index 967ddbc..6c0f265 100644 --- a/tests/unit/test_tools.py +++ b/tests/unit/test_tools.py @@ -195,7 +195,7 @@ async def test_get_server_info_is_callable_from_an_mcp_client() -> None: async def test_list_databases_primary_only(monkeypatch: pytest.MonkeyPatch) -> None: """With no secondaries, list_databases reports just the primary.""" - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver db = Database(_SETTINGS, pool=FakePool()) # type: ignore[arg-type] @@ -220,7 +220,7 @@ async def execute_query(self, *a: Any, **k: Any) -> list[SqlDriver.RowResult]: async def test_list_databases_lists_secondaries(monkeypatch: pytest.MonkeyPatch) -> None: - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver settings = load_settings( { @@ -253,7 +253,7 @@ async def execute_query(self, *a: Any, **k: Any) -> list[SqlDriver.RowResult]: async def test_list_databases_surfaces_unreachable_secondary(monkeypatch: pytest.MonkeyPatch) -> None: - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver settings = load_settings( { @@ -293,7 +293,7 @@ def _driver(database_id: str | None = None) -> Any: async def test_read_tool_threads_database_param_to_driver(monkeypatch: pytest.MonkeyPatch) -> None: """A read tool's ``database`` arg must select the secondary's driver.""" - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver settings = load_settings( { diff --git a/tests/unit/test_vector_ops.py b/tests/unit/test_vector_ops.py index 73c78f9..efab1fd 100644 --- a/tests/unit/test_vector_ops.py +++ b/tests/unit/test_vector_ops.py @@ -1638,7 +1638,7 @@ async def execute_query( params: list[object] | None = None, force_readonly: bool = False, ) -> list[object]: - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver self.calls.append((sql, list(params or []))) if "pg_extension" in sql: @@ -1692,7 +1692,7 @@ async def execute_query( params: list[object] | None = None, force_readonly: bool = False, ) -> list[object]: - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver self.calls.append((sql, list(params or []))) if "pg_extension" in sql: @@ -1755,7 +1755,7 @@ async def execute_query( params: list[object] | None = None, force_readonly: bool = False, ) -> list[object]: - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver self.calls.append((sql, list(params or []))) if "pg_extension" in sql: diff --git a/tests/unit/test_wait_for_lsn.py b/tests/unit/test_wait_for_lsn.py index 6175055..2f20106 100644 --- a/tests/unit/test_wait_for_lsn.py +++ b/tests/unit/test_wait_for_lsn.py @@ -105,7 +105,7 @@ def __init__(self, *, ver_num: int = 190001, ver: str = "19beta1", fail_with: st self.calls: list[tuple[str, object, bool]] = [] async def execute_query(self, query, params=None, force_readonly=False): # type: ignore[no-untyped-def] - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver self.calls.append((query, params, force_readonly)) if "current_setting" in query: @@ -159,7 +159,7 @@ async def test_wait_for_lsn_timeout_detected_via_sqlstate_57014() -> None: driver._fail_with = None async def fail_with_sqlstate(query, params=None, force_readonly=False): # type: ignore[no-untyped-def] - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver driver.calls.append((query, params, force_readonly)) if "current_setting" in query: diff --git a/tests/unit/test_warehousepg.py b/tests/unit/test_warehousepg.py index ed39c88..41a988c 100644 --- a/tests/unit/test_warehousepg.py +++ b/tests/unit/test_warehousepg.py @@ -151,7 +151,7 @@ async def execute_query( self.call_index += 1 if self.call_index == 1: # Version probe — succeed with a WarehousePG banner. - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver return [SqlDriver.RowResult(cells={"version": _WAREHOUSEPG_VERSION})] raise RuntimeError("permission denied for view gp_segment_configuration") @@ -179,7 +179,7 @@ async def execute_query( ) -> list[Any]: del params, force_readonly self.call_index += 1 - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver if self.call_index == 1: return [SqlDriver.RowResult(cells={"version": _WAREHOUSEPG_VERSION})] diff --git a/tests/unit/test_warehousepg_advisors.py b/tests/unit/test_warehousepg_advisors.py index 44368ca..246b4f4 100644 --- a/tests/unit/test_warehousepg_advisors.py +++ b/tests/unit/test_warehousepg_advisors.py @@ -155,7 +155,7 @@ async def execute_query( self, query: str, params: list[Any] | None = None, force_readonly: bool = False ) -> list[Any]: del params, force_readonly - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver self.call_index += 1 if self.call_index == 1: diff --git a/tests/unit/test_warehousepg_reads.py b/tests/unit/test_warehousepg_reads.py index b863d81..2a12c03 100644 --- a/tests/unit/test_warehousepg_reads.py +++ b/tests/unit/test_warehousepg_reads.py @@ -361,7 +361,7 @@ async def execute_query( ) -> list[Any]: del force_readonly self.call_index += 1 - from mcpg._vendor.sql import SqlDriver + from mcpg.sql import SqlDriver # Status probe — 3 calls (version + to_regclass + segments). if self.call_index == 1: diff --git a/tests/vendor/Dockerfile.postgres-hypopg b/tests/vendor/Dockerfile.postgres-hypopg deleted file mode 100644 index 3741b0e..0000000 --- a/tests/vendor/Dockerfile.postgres-hypopg +++ /dev/null @@ -1,21 +0,0 @@ -ARG PG_VERSION=15 -FROM postgres:${PG_VERSION} - -# Install build dependencies -RUN apt-get update && apt-get install -y \ - build-essential \ - git \ - postgresql-server-dev-${PG_MAJOR} \ - && rm -rf /var/lib/apt/lists/* - -# Clone and build HypoPG -RUN git clone --depth 1 https://github.com/HypoPG/hypopg.git /tmp/hypopg \ - && cd /tmp/hypopg \ - && make \ - && make install \ - && cd / \ - && rm -rf /tmp/hypopg - -# Add initialization script to create extensions in template1 -RUN echo "CREATE EXTENSION IF NOT EXISTS pg_stat_statements; CREATE EXTENSION IF NOT EXISTS hypopg;" \ - > /docker-entrypoint-initdb.d/00-create-extensions.sql diff --git a/tests/vendor/sql/test_obfuscate_password.py b/tests/vendor/sql/test_obfuscate_password.py deleted file mode 100644 index 08fbd7a..0000000 --- a/tests/vendor/sql/test_obfuscate_password.py +++ /dev/null @@ -1,104 +0,0 @@ -from mcpg._vendor.sql import obfuscate_password - - -def test_obfuscate_none_or_empty(): - """Test that None or empty strings are handled correctly.""" - assert obfuscate_password("") == "" - assert obfuscate_password(None) is None - - -def test_obfuscate_postgresql_url(): - """Test obfuscation of regular PostgreSQL connection URLs.""" - # Standard URL - url = "postgresql://user:secret@localhost:5432/mydatabase" - result = obfuscate_password(url) - assert result is not None - assert "secret" not in result - assert "****" in result - assert result == "postgresql://user:****@localhost:5432/mydatabase" - - # URL with special characters in password - url = "postgresql://user:p@$$w0rd@localhost:5432/mydatabase" - result = obfuscate_password(url) - assert result is not None - assert "p@$$w0rd" not in result - assert "****" in result - - # URL with query parameters - url = "postgresql://user:secret@localhost:5432/mydatabase?sslmode=require" - result = obfuscate_password(url) - assert result is not None - assert "secret" not in result - assert "?sslmode=require" in result - - -def test_obfuscate_in_error_message(): - """Test obfuscation of URLs within error messages.""" - error_msg = ( - "Failed to connect: could not connect to server: Connection refused. Is the server " - "running on host 'localhost' (127.0.0.1) and accepting TCP/IP connections on port 5432? " - "connection string: postgresql://admin:topsecret@localhost:5432/mydb" - ) - obfuscated = obfuscate_password(error_msg) - assert obfuscated is not None - assert "topsecret" not in obfuscated - assert "****" in obfuscated - assert "postgresql://admin:****@localhost:5432/mydb" in obfuscated - - -def test_obfuscate_connection_params(): - """Test obfuscation of connection parameters.""" - # Key=value format - conn_string = "host=localhost port=5432 dbname=mydb user=admin password=secret123" - obfuscated = obfuscate_password(conn_string) - assert obfuscated is not None - assert "secret123" not in obfuscated - assert "password=****" in obfuscated - - # Connection in Python code with single quotes - code_snippet = """conn = psycopg.connect("host=localhost dbname=mydb user=postgres password='my$3cret!'")""" - obfuscated = obfuscate_password(code_snippet) - assert obfuscated is not None - assert "my$3cret!" not in obfuscated - assert "password='****'" in obfuscated - - -def test_obfuscate_multiple_passwords(): - """Test obfuscation of multiple passwords in the same string.""" - text = """ - Primary DB: postgresql://user1:password1@host1:5432/db1 - Secondary DB: postgresql://user2:password2@host2:5432/db2 - """ - obfuscated = obfuscate_password(text) - assert obfuscated is not None - assert "password1" not in obfuscated - assert "password2" not in obfuscated - assert "user1:****@" in obfuscated - assert "user2:****@" in obfuscated - - -def test_obfuscate_no_sensitive_data(): - """Test that strings without sensitive data are unchanged.""" - text = "This is a normal string with no passwords." - assert obfuscate_password(text) == text - - # URL without password - url = "http://example.com/path" - assert obfuscate_password(url) == url - - -def test_obfuscate_dsn_format(): - """Test obfuscation of DSN format passwords.""" - # Single quotes - dsn = "host='localhost' user='postgres' password='supersecret' dbname='testdb'" - obfuscated = obfuscate_password(dsn) - assert obfuscated is not None - assert "supersecret" not in obfuscated - assert "password='****'" in obfuscated - - # Double quotes - dsn = 'host="localhost" user="postgres" password="supersecret" dbname="testdb"' - obfuscated = obfuscate_password(dsn) - assert obfuscated is not None - assert "supersecret" not in obfuscated - assert 'password="****"' in obfuscated diff --git a/tests/vendor/sql/test_safe_sql.py b/tests/vendor/sql/test_safe_sql.py deleted file mode 100644 index cbf5132..0000000 --- a/tests/vendor/sql/test_safe_sql.py +++ /dev/null @@ -1,760 +0,0 @@ -from unittest.mock import AsyncMock -from unittest.mock import Mock -from unittest.mock import call - -import pytest -import pytest_asyncio -from psycopg.sql import SQL -from psycopg.sql import Literal - -from mcpg._vendor.sql import SafeSqlDriver -from mcpg._vendor.sql import SqlDriver - - -@pytest_asyncio.fixture -async def mock_sql_driver(): - driver = Mock(spec=SqlDriver) - driver.execute_query = AsyncMock(return_value=[]) - return driver - - -@pytest_asyncio.fixture -async def safe_driver(mock_sql_driver): - return SafeSqlDriver(mock_sql_driver) - - -@pytest.mark.asyncio -async def test_select_statement(safe_driver, mock_sql_driver): - """Test that simple SELECT statements are allowed""" - query = "SELECT * FROM users WHERE age > 18" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_update_statement(safe_driver): - """Test that UPDATE statements are blocked""" - query = "UPDATE users SET status = 'active' WHERE id = 1" - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_select_with_join(safe_driver, mock_sql_driver): - """Test that SELECT with JOIN is allowed""" - query = """ - SELECT users.name, orders.order_date - FROM users - INNER JOIN orders ON users.id = orders.user_id - WHERE orders.status = 'pending' - """ - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_show_variable(safe_driver, mock_sql_driver): - """Test that SHOW statements are allowed""" - query = "SHOW search_path" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_set_variable(safe_driver): - """Test that SET statements are blocked""" - query = "SET search_path TO public" - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_select_with_arithmetic(safe_driver, mock_sql_driver): - """Test that SELECT with arithmetic expressions is allowed""" - query = "SELECT id, price * quantity as total FROM orders" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_select_current_user(safe_driver, mock_sql_driver): - """Test that SELECT current_user is allowed""" - query = "SELECT current_user" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_drop_table(safe_driver): - """Test that DROP TABLE statements are blocked""" - query = "DROP TABLE users" - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_delete_from_table(safe_driver): - """Test that DELETE FROM statements are blocked""" - query = "DELETE FROM users WHERE status = 'inactive'" - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_select_with_subquery(safe_driver, mock_sql_driver): - """Test that SELECT with subqueries is allowed""" - query = """ - SELECT name FROM users - WHERE id IN (SELECT user_id FROM orders WHERE total > 1000) - """ - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_select_with_malicious_comment(safe_driver): - """Test that SQL injection via comments is blocked""" - query = """ - SELECT * FROM users; DROP TABLE users; - """ - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_select_with_union(safe_driver, mock_sql_driver): - """Test that UNION queries are allowed""" - query = """ - SELECT id, name FROM users - UNION - SELECT NULL, concat(table_name) FROM information_schema.tables - """ - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_select_into(safe_driver): - """Test that SELECT INTO statements are blocked""" - query = """ - SELECT id, name - INTO new_table - FROM users - """ - with pytest.raises(ValueError, match="Error validating query"): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_select_for_update(safe_driver): - """Test that SELECT FOR UPDATE statements are blocked""" - query = """ - SELECT id, name - FROM users - FOR UPDATE - """ - with pytest.raises(ValueError, match="Error validating query"): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_select_with_locking_clause(safe_driver): - """Test that SELECT with explicit locking clauses is blocked""" - query = """ - SELECT id, name - FROM users - FOR SHARE NOWAIT - """ - with pytest.raises(ValueError, match="Error validating query"): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_select_with_commit(safe_driver): - """Test that statements containing COMMIT are blocked""" - query = """ - SELECT id FROM users; - COMMIT; - SELECT name FROM users; - """ - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_explain_plan(safe_driver, mock_sql_driver): - """Test that EXPLAIN (without ANALYZE) works with bind variables""" - query = """ - EXPLAIN (FORMAT JSON) - SELECT id, name - FROM users - WHERE age > $1 AND status = $2 - """ - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_explain_analyze_blocked(safe_driver): - """Test that EXPLAIN ANALYZE is blocked""" - query = """ - EXPLAIN ANALYZE - SELECT id, name FROM users - """ - with pytest.raises(ValueError, match="Error validating query"): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_begin_transaction_blocked(safe_driver): - """Test that transaction blocks are blocked""" - query = """ - BEGIN; - SELECT id, name FROM users; - """ - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_invalid_sql_syntax(safe_driver): - """Test that queries with invalid SQL syntax are blocked""" - query1 = "SELECT * FRMO users;" - with pytest.raises(ValueError, match="Failed to parse SQL statement"): - await safe_driver.execute_query(query1) - - query2 = "SELECT * FROM users INNER JOON posts ON users.id = posts.user_id;" - with pytest.raises(ValueError, match="Failed to parse SQL statement"): - await safe_driver.execute_query(query2) - - query3 = "SELECT * FROM users WHERE (age > 21 AND active = true;" - with pytest.raises(ValueError, match="Failed to parse SQL statement"): - await safe_driver.execute_query(query3) - - -@pytest.mark.asyncio -async def test_create_index_blocked(safe_driver): - """Test that CREATE INDEX statements are blocked""" - query = """ - CREATE INDEX idx_user_email ON users(email); - """ - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_drop_index_blocked(safe_driver): - """Test that DROP INDEX statements are blocked""" - query = """ - DROP INDEX idx_user_email; - """ - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_create_table_blocked(safe_driver): - """Test that CREATE TABLE statements are blocked""" - query = """ - CREATE TABLE test_table ( - id SERIAL PRIMARY KEY, - name VARCHAR(100) NOT NULL, - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP - ); - """ - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_create_table_as_blocked(safe_driver): - """Test that CREATE TABLE AS statements are blocked""" - query = """ - CREATE TABLE user_backup AS - SELECT * FROM users; - """ - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_create_extension_blocked(safe_driver): - """Test that CREATE EXTENSION statements are blocked""" - query = """ - CREATE EXTENSION pg_hack; - """ - with pytest.raises(ValueError, match="Error validating query"): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_drop_extension_blocked(safe_driver): - """Test that DROP EXTENSION statements are blocked""" - query = """ - DROP EXTENSION pg_stat_statements; - """ - with pytest.raises( - ValueError, - match="Error validating query", - ): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_complex_index_metadata_select(safe_driver, mock_sql_driver): - """Test that complex SELECT queries for index metadata are allowed""" - query = """SELECT indexrelid::regclass AS index_name, array_agg(attname) AS columns, - indisunique, indisprimary - FROM pg_index i - JOIN pg_attribute a ON a.attnum = ANY(i.indkey) - WHERE i.indrelid IN (SELECT oid FROM pg_class WHERE relnamespace = - (SELECT oid FROM pg_namespace WHERE nspname = 'public')) - GROUP BY indexrelid, indisunique, indisprimary - HAVING COUNT(array_agg(attname)) > 1""" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_allowed_functions(safe_driver): - """Tests that allow functions (especially the ones that are newly added)""" - query = """ - SELECT pg_relation_filenode('foo'); - """ - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_disallowed_functions(safe_driver): - """Test that disallowed functions are blocked""" - queries = [ - "SELECT pg_sleep(1);", - "SELECT pg_read_file('/etc/passwd');", - "SELECT lo_import('/etc/passwd');", - ] - - for query in queries: - with pytest.raises(ValueError, match="Error validating query"): - await safe_driver.execute_query(query) - - -@pytest.mark.asyncio -async def test_session_info_functions(safe_driver, mock_sql_driver): - """Test that session info functions are allowed""" - query = "SELECT current_user, current_database(), version()" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_blocking_pids_functions(safe_driver, mock_sql_driver): - """Test that blocking pids functions are allowed""" - query = "SELECT pg_blocking_pids(1234)" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_logfile_functions(safe_driver, mock_sql_driver): - """Test that logfile functions are allowed""" - query = "SELECT pg_current_logfile()" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_complex_session_info_queries(safe_driver, mock_sql_driver): - """Test that complex session info queries are allowed""" - query = """ - SELECT current_user, current_database(), version(), - pg_backend_pid(), pg_blocking_pids(pg_backend_pid()) - """ - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_security_privilege_functions(safe_driver, mock_sql_driver): - """Test that security privilege functions are allowed""" - query = "SELECT has_table_privilege('user', 'table', 'SELECT')" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_complex_security_privilege_queries(safe_driver, mock_sql_driver): - """Test more complex queries using security privilege functions""" - queries = [ - """ - SELECT tablename, has_table_privilege(current_user, tablename, 'SELECT') as can_select, - has_table_privilege(current_user, tablename, 'UPDATE') as can_update - FROM pg_tables - WHERE schemaname = 'public' - ORDER BY tablename - """, - """ - SELECT proname, has_function_privilege(current_user, p.oid, 'EXECUTE') as can_execute - FROM pg_proc p - JOIN pg_namespace n ON p.pronamespace = n.oid - WHERE n.nspname = 'public' - """, - """ - SELECT c.relname, a.attname, - has_column_privilege(current_user, c.oid, a.attnum, 'SELECT') as can_select, - has_column_privilege(current_user, c.oid, a.attnum, 'UPDATE') as can_update - FROM pg_class c - JOIN pg_attribute a ON c.oid = a.attrelid - WHERE c.relname = 'mytable' - AND a.attnum > 0 - AND NOT a.attisdropped - ORDER BY a.attnum - """, - ] - - for query in queries: - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_security_privilege_functions_with_subqueries(safe_driver, mock_sql_driver): - """Test security privilege functions used within subqueries""" - queries = [ - """ - SELECT nspname - FROM pg_namespace - WHERE has_schema_privilege(current_user, oid, 'USAGE') - AND nspname NOT LIKE 'pg_%' - ORDER BY nspname - """, - """ - SELECT t.tablename - FROM pg_tables t - WHERE EXISTS ( - SELECT 1 - FROM information_schema.columns c - WHERE c.table_name = t.tablename - AND has_column_privilege(current_user, t.tablename, c.column_name, 'SELECT') - ) - AND t.schemaname = 'public' - """, - ] - - for query in queries: - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.parametrize("operator", ["LIKE", "ILIKE"]) -@pytest.mark.asyncio -async def test_like_patterns(safe_driver, mock_sql_driver, operator): - """Test that LIKE/ILIKE patterns are only allowed if they start or end with %, but not both or in middle""" - queries = [ - f"SELECT name FROM users WHERE name {operator} '%smith'", - f"SELECT name FROM users WHERE name {operator} 'john%'", - f"SELECT name FROM users WHERE name {operator} 'jo%hn'", - f"SELECT name FROM users WHERE name {operator} '%jo%hn%'", - f"SELECT name FROM users WHERE name {operator} '%john%'", - ] - - for query in queries: - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_datetime_functions(safe_driver, mock_sql_driver): - """Test that date/time functions are allowed""" - queries = [ - "SELECT age(timestamp '2001-04-10', timestamp '1957-06-13')", - "SELECT clock_timestamp()", - "SELECT current_date", - "SELECT current_time", - "SELECT current_timestamp", - "SELECT date_part('year', timestamp '2001-02-16')", - "SELECT date_trunc('hour', timestamp '2001-02-16 20:38:40')", - "SELECT extract(year from timestamp '2001-02-16 20:38:40')", - "SELECT isfinite(timestamp '2001-02-16')", - "SELECT justify_days(interval '35 days')", - "SELECT justify_hours(interval '27 hours')", - "SELECT justify_interval(interval '1 year -1 hour')", - "SELECT localtime", - "SELECT localtimestamp", - "SELECT make_date(2013, 7, 15)", - "SELECT make_interval(years := 1)", - "SELECT make_time(8, 15, 23.5)", - "SELECT make_timestamp(2013, 7, 15, 8, 15, 23.5)", - "SELECT make_timestamptz(2013, 7, 15, 8, 15, 23.5)", - "SELECT now()", - "SELECT statement_timestamp()", - "SELECT timeofday()", - "SELECT transaction_timestamp()", - ] - - for query in queries: - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_type_conversion_functions(safe_driver, mock_sql_driver): - """Test that type conversion functions are allowed""" - queries = [ - "SELECT CAST('100' AS integer)", - "SELECT '100'::integer", - "SELECT text '100'", - "SELECT bool 'true'", - "SELECT int2 '100'", - "SELECT int4 '100'", - "SELECT int8 '100'", - "SELECT float4 '100.0'", - "SELECT float8 '100.0'", - "SELECT numeric '100.0'", - "SELECT date '2001-10-05'", - "SELECT time '04:05:06.789'", - "SELECT timetz '04:05:06.789-08'", - "SELECT timestamp '2001-10-05 04:05:06.789'", - "SELECT timestamptz '2001-10-05 04:05:06.789-08'", - "SELECT interval '1 year 2 months 3 days'", - ] - - for query in queries: - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_regexp_functions(safe_driver, mock_sql_driver): - """Test that regexp functions are allowed""" - query = "SELECT regexp_replace('Hello World', 'World', 'PostgreSQL')" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_complex_type_conversion_queries(safe_driver, mock_sql_driver): - """Test that complex type conversion queries are allowed""" - query = """ - SELECT to_char(current_timestamp, 'YYYY-MM-DD'), - to_date('2023-01-01'), - to_timestamp('2023-01-01 12:00:00'), - to_number('123.45', '999.99') - """ - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_network_functions(safe_driver, mock_sql_driver): - """Test that network functions are allowed""" - query = "SELECT inet_client_addr(), inet_client_port()" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_network_functions_in_complex_queries(safe_driver, mock_sql_driver): - """Test that network functions in complex queries are allowed""" - query = """ - SELECT inet_client_addr() as client_ip, - inet_client_port() as client_port, - inet_server_addr() as server_ip, - inet_server_port() as server_port - """ - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_notification_and_server_functions(safe_driver, mock_sql_driver): - """Test that notification and server functions are allowed""" - query = "SELECT pg_listening_channels(), pg_postmaster_start_time()" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_minmax_expressions(safe_driver, mock_sql_driver): - """Test that minmax expressions are allowed""" - query = "SELECT GREATEST(1, 2, 3), LEAST(1, 2, 3)" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_row_expressions(safe_driver, mock_sql_driver): - """Test that row expressions are allowed""" - query = "SELECT ROW(1, 2, 3) = ROW(1, 2, 3)" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_extension_check_query(safe_driver, mock_sql_driver): - """Test that extension check queries are allowed""" - query = "SELECT extname, extversion FROM pg_extension WHERE extname = 'hypopg'" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_create_extension_query(safe_driver, mock_sql_driver): - """Test that CREATE EXTENSION queries are allowed""" - query = "CREATE EXTENSION IF NOT EXISTS hypopg" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_hypopg_create_index_query(safe_driver, mock_sql_driver): - """Test that hypopg create index queries are allowed""" - query = "SELECT * FROM hypopg_create_index('CREATE INDEX idx ON users(id)')" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_hypopg_reset_query(safe_driver, mock_sql_driver): - """Test that hypopg reset queries are allowed""" - query = "SELECT * FROM hypopg_reset()" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_hypopg_list_indexes_query(safe_driver, mock_sql_driver): - """Test that hypopg list indexes queries are allowed""" - query = "SELECT * FROM hypopg_list_indexes()" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_pg_stat_statements_query(safe_driver, mock_sql_driver): - """Test that pg_stat_statements queries are allowed""" - query = "SELECT * FROM pg_stat_statements ORDER BY calls DESC LIMIT 10" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_pg_indexes_query(safe_driver, mock_sql_driver): - """Test that pg_indexes queries are allowed""" - query = "SELECT * FROM pg_indexes WHERE schemaname = 'public'" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_pg_stats_query(safe_driver, mock_sql_driver): - """Test that pg_stats queries are allowed""" - query = "SELECT * FROM pg_stats WHERE schemaname = 'public'" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_explain_query(safe_driver, mock_sql_driver): - """Test that explain queries are allowed""" - query = "EXPLAIN SELECT * FROM users" - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_sql_driver_parameter_format(safe_driver, mock_sql_driver): - """Test query with SQL parameters through the DatabaseTuningAdvisor.""" - query_template = """ - SELECT queryid, query, calls, total_exec_time/calls as avg_exec_time - FROM pg_stat_statements - WHERE calls >= {} - AND total_exec_time/calls >= {} - ORDER BY total_exec_time DESC - LIMIT {} - """ - - min_calls = 50 - min_avg_time = 5.0 - limit = 100 - - formatted_query = SQL(query_template).format(Literal(min_calls), Literal(min_avg_time), Literal(limit)).as_string() - - await safe_driver.execute_query(formatted_query) - mock_sql_driver.execute_query.assert_awaited_with("/* crystaldba */ " + formatted_query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_multiple_queries(safe_driver, mock_sql_driver): - """Test that multiple queries are handled correctly""" - query1 = "SELECT * FROM users" - query2 = "SELECT * FROM orders" - await safe_driver.execute_query(query1) - await safe_driver.execute_query(query2) - mock_sql_driver.execute_query.assert_has_awaits( - [ - call("/* crystaldba */ " + query1, params=None, force_readonly=True), - call("/* crystaldba */ " + query2, params=None, force_readonly=True), - ] - ) - - -@pytest.mark.asyncio -async def test_query_with_comments(safe_driver, mock_sql_driver): - """Test that queries with comments are handled correctly""" - query = """ - -- Get user information - SELECT id, name, email - FROM users - WHERE status = 'active' - -- Only get active users - """ - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) - - -@pytest.mark.asyncio -async def test_query_with_whitespace(safe_driver, mock_sql_driver): - """Test that queries with whitespace are handled correctly""" - query = """ - SELECT id, - name, - email - FROM users - WHERE status = 'active' - ORDER BY name - """ - await safe_driver.execute_query(query) - mock_sql_driver.execute_query.assert_awaited_once_with("/* crystaldba */ " + query, params=None, force_readonly=True) diff --git a/tests/vendor/sql/test_sql_driver.py b/tests/vendor/sql/test_sql_driver.py deleted file mode 100644 index 9733c3a..0000000 --- a/tests/vendor/sql/test_sql_driver.py +++ /dev/null @@ -1,367 +0,0 @@ -# ruff: noqa: B017 -from unittest.mock import AsyncMock -from unittest.mock import MagicMock -from unittest.mock import call -from unittest.mock import patch - -import pytest - -from mcpg._vendor.sql import DbConnPool -from mcpg._vendor.sql import SqlDriver - - -class AsyncContextManagerMock(AsyncMock): - """A better mock for async context managers""" - - async def __aenter__(self): - return self.aenter - - async def __aexit__(self, exc_type, exc_val, exc_tb): - pass - - -@pytest.fixture -def mock_connection(): - """Create a mock for the database connection.""" - connection = MagicMock() - cursor = AsyncContextManagerMock() - - # Make the cursor behave like an async context manager - cursor.aenter = cursor - - # Configure cursor context manager - cursor_cm = AsyncContextManagerMock() - cursor_cm.aenter = cursor - connection.cursor.return_value = cursor_cm - - # Configure cursor.description to return a value (indicating results) - cursor.description = ["column1", "column2"] - - # Configure fetchall to return some mock data - cursor.fetchall.return_value = [ - {"id": 1, "name": "test1"}, - {"id": 2, "name": "test2"}, - ] - - return connection, cursor - - -@pytest.fixture -def mock_db_pool(): - """Create a mock for DbConnPool with a mock connection.""" - # Create the pool - pool = MagicMock() - - # Create connection that returns from async context manager - connection = AsyncContextManagerMock() - connection.aenter = connection - - # Create cursor that returns from async context manager - cursor = AsyncContextManagerMock() - cursor.aenter = cursor - - # Setup connection to return cursor - cursor_cm = AsyncContextManagerMock() - cursor_cm.aenter = cursor - connection.cursor.return_value = cursor_cm - - # Configure cursor.description - cursor.description = ["column1", "column2"] - - # Configure fetchall to return some mock data - cursor.fetchall.return_value = [ - {"id": 1, "name": "test1"}, - {"id": 2, "name": "test2"}, - ] - - # Create connection context manager - conn_cm = AsyncContextManagerMock() - conn_cm.aenter = connection - - # Setup pool.connection() to return the connection context manager - pool.connection.return_value = conn_cm - - # Create mock for DbConnPool - db_pool = MagicMock(spec=DbConnPool) - db_pool.pool_connect.return_value = pool - db_pool._is_valid = True - - return db_pool, connection, cursor - - -@pytest.mark.asyncio -async def test_execute_query_readonly_transaction(mock_connection): - """Test execute_query with read-only transaction.""" - connection, cursor = mock_connection - - # Create SqlDriver with a connection - driver = SqlDriver(conn=connection) - - # Create a mock implementation of _execute_with_connection - async def mock_impl(connection, query, params, force_readonly): - # Simulate transaction - await cursor.execute("BEGIN TRANSACTION READ ONLY" if force_readonly else "BEGIN TRANSACTION") - - # Execute the query - if params: - await cursor.execute(query, params) - else: - await cursor.execute(query) - - # Get results - rows = await cursor.fetchall() - - # End transaction - if force_readonly: - await cursor.execute("ROLLBACK") - else: - await cursor.execute("COMMIT") - - # Return results - return [SqlDriver.RowResult(cells=dict(row)) for row in rows] - - # Must match the parameter names from the original method - driver._execute_with_connection = mock_impl # type: ignore - - # Execute a read-only query - result = await driver._execute_with_connection( # type: ignore - connection, "SELECT * FROM test", None, force_readonly=True - ) - - # Verify transaction management - assert cursor.execute.call_count >= 3 - assert call("BEGIN TRANSACTION READ ONLY") in cursor.execute.call_args_list - assert call("ROLLBACK") in cursor.execute.call_args_list - - # Verify results were processed correctly - assert result is not None - assert len(result) == 2 - assert result[0].cells["id"] == 1 - assert result[1].cells["name"] == "test2" - - -@pytest.mark.asyncio -async def test_execute_query_writeable_transaction(mock_connection): - """Test execute_query with writeable transaction.""" - connection, cursor = mock_connection - - # Create SqlDriver with a connection - driver = SqlDriver(conn=connection) - - # Create a mock implementation of _execute_with_connection - async def mock_impl(connection, query, params, force_readonly): - # Simulate transaction - await cursor.execute("BEGIN TRANSACTION READ ONLY" if force_readonly else "BEGIN TRANSACTION") - - # Execute the query - if params: - await cursor.execute(query, params) - else: - await cursor.execute(query) - - # Get results - rows = await cursor.fetchall() - - # End transaction - if force_readonly: - await cursor.execute("ROLLBACK") - else: - await cursor.execute("COMMIT") - - # Return results - return [SqlDriver.RowResult(cells=dict(row)) for row in rows] - - # Must match the parameter names from the original method - driver._execute_with_connection = mock_impl # type: ignore - - # Execute a writeable query - result = await driver._execute_with_connection( # type: ignore - connection, "UPDATE test SET name = 'updated'", None, force_readonly=False - ) - - # Verify transaction management - assert call("COMMIT") in cursor.execute.call_args_list - - # Verify results were processed correctly - assert result is not None - - -@pytest.mark.asyncio -async def test_execute_query_error_handling(mock_connection): - """Test execute_query error handling.""" - connection, cursor = mock_connection - - # Configure cursor.execute to raise an exception on the second call - cursor.execute.side_effect = [None, Exception("Query execution failed")] - - # Create SqlDriver with a connection - driver = SqlDriver(conn=connection) - - # Create mock function that raises exception - async def mock_execute_error(connection, query, params, force_readonly): - raise Exception("Query execution failed") - - driver._execute_with_connection = mock_execute_error # type: ignore - - # Execute a query that will fail - with pytest.raises(Exception) as excinfo: - await driver._execute_with_connection( # type: ignore - connection, "SELECT * FROM nonexistent", None, force_readonly=True - ) - - # Check the error message - assert "Query execution failed" in str(excinfo.value) - - -@pytest.mark.asyncio -async def test_execute_query_no_results(mock_connection): - """Test execute_query with no results returned.""" - connection, cursor = mock_connection - - # Configure cursor.description to return None (indicating no results) - cursor.description = None - - # Create SqlDriver with a connection - driver = SqlDriver(conn=connection) - - # Create a mock implementation of _execute_with_connection - async def mock_impl(connection, query, params, force_readonly): - # Simulate transaction - await cursor.execute("BEGIN TRANSACTION READ ONLY" if force_readonly else "BEGIN TRANSACTION") - - # Execute the query - if params: - await cursor.execute(query, params) - else: - await cursor.execute(query) - - # End transaction - if force_readonly: - await cursor.execute("ROLLBACK") - else: - await cursor.execute("COMMIT") - - # Return None for no results - return None - - # Must match the parameter names from the original method - driver._execute_with_connection = mock_impl # type: ignore - - # Execute a query that returns no results - result = await driver._execute_with_connection( # type: ignore - connection, "DELETE FROM test", None, force_readonly=False - ) - - # Verify result is None for no-result queries - assert result is None - assert call("COMMIT") in cursor.execute.call_args_list - - -@pytest.mark.asyncio -async def test_execute_query_with_params(mock_connection): - """Test execute_query with parameters.""" - connection, cursor = mock_connection - - # Create SqlDriver with a connection - driver = SqlDriver(conn=connection) - - # Create a mock implementation of _execute_with_connection - async def mock_impl(connection, query, params, force_readonly): - # Simulate transaction - await cursor.execute("BEGIN TRANSACTION READ ONLY" if force_readonly else "BEGIN TRANSACTION") - - # Execute the query with parameters - if params: - await cursor.execute(query, params) - else: - await cursor.execute(query) - - # Get results - rows = await cursor.fetchall() - - # End transaction - if force_readonly: - await cursor.execute("ROLLBACK") - else: - await cursor.execute("COMMIT") - - # Return results - return [SqlDriver.RowResult(cells=dict(row)) for row in rows] - - # Must match the parameter names from the original method - driver._execute_with_connection = mock_impl # type: ignore - - # Execute a query with parameters - await driver._execute_with_connection( # type: ignore - connection, "SELECT * FROM test WHERE id = %s", [1], force_readonly=True - ) - - # Verify parameters were passed correctly - assert call("SELECT * FROM test WHERE id = %s", [1]) in cursor.execute.call_args_list - - -@pytest.mark.asyncio -async def test_execute_query_from_pool(mock_db_pool): - """Test execute_query using a connection from a pool.""" - db_pool, connection, cursor = mock_db_pool - - # Create a mock execute function - async def mock_pool_execute(*args, **kwargs): - return [ - SqlDriver.RowResult(cells={"id": 1, "name": "test1"}), - SqlDriver.RowResult(cells={"id": 2, "name": "test2"}), - ] - - # Create SqlDriver with the mocked pool - driver = SqlDriver(conn=db_pool) - driver.execute_query = mock_pool_execute - - # Execute a query - result = await driver.execute_query("SELECT * FROM test") - - # Verify results were processed correctly - assert result is not None - assert len(result) == 2 - assert result[0].cells["id"] == 1 - assert result[1].cells["name"] == "test2" - - -@pytest.mark.asyncio -async def test_connection_error_marks_pool_invalid(mock_db_pool): - """Test that connection errors mark the pool as invalid.""" - db_pool, connection, cursor = mock_db_pool - - # Configure pool_connect to raise an exception - db_pool.pool_connect.side_effect = Exception("Connection failed") - - # Create SqlDriver with the mocked pool - driver = SqlDriver(conn=db_pool) - - # Execute a query that will fail due to connection error - with pytest.raises(Exception): - await driver.execute_query("SELECT * FROM test") - - # Make pool invalid manually (since we're bypassing the actual method) - db_pool._is_valid = False - db_pool._last_error = "Connection failed" - - # Verify pool was marked as invalid - assert db_pool._is_valid is False - assert isinstance(db_pool._last_error, str) - - -@pytest.mark.asyncio -async def test_engine_url_connection(): - """Test connecting with engine_url instead of connection object.""" - db_pool = MagicMock(spec=DbConnPool) - - with patch("mcpg._vendor.sql.DbConnPool", return_value=db_pool): - # Create SqlDriver with engine_url - driver = SqlDriver(engine_url="postgresql://user:pass@localhost/db") - - # Call connect to create mock pool - driver.connect() - - # Verify driver state - assert driver.is_pool is True - assert driver.conn is not None diff --git a/tests/vendor/utils.py b/tests/vendor/utils.py deleted file mode 100644 index 7780cc2..0000000 --- a/tests/vendor/utils.py +++ /dev/null @@ -1,155 +0,0 @@ -import logging -import os -import time -from pathlib import Path -from typing import Generator -from typing import Tuple - -import docker -import pytest -from docker import errors as docker_errors - -logger = logging.getLogger(__name__) - - -def create_postgres_container(version: str) -> Generator[Tuple[str, str], None, None]: - """Create a PostgreSQL container of specified version and return its connection string.""" - try: - client = docker.from_env() - client.ping() - except (docker_errors.DockerException, ConnectionError): - pytest.skip("Docker is not available") - - # Extract PostgreSQL version number - pg_version = version.split(":")[1] if ":" in version else version - - # Define custom image name with HypoPG - custom_image_name = f"postgres-hypopg:{pg_version}" - - container_name = f"postgres-crystal-test-{version.replace(':', '_')}-{os.urandom(4).hex()}" - current_dir = Path(__file__).parent.absolute() - - logger.info(f"Setting up PostgreSQL {pg_version} with HypoPG") - - # Build custom Docker image with HypoPG if it doesn't exist - try: - # Check if custom image already exists - client.images.get(custom_image_name) - logger.info(f"Using existing Docker image: {custom_image_name}") - except docker_errors.ImageNotFound: - # Build the custom image - logger.info(f"Building custom Docker image: {custom_image_name}") - try: - dockerfile_path = current_dir / "Dockerfile.postgres-hypopg" - if not dockerfile_path.exists(): - logger.error(f"Dockerfile not found at {dockerfile_path}") - pytest.skip(f"Required Dockerfile not found: {dockerfile_path}") - - # Build the image - client.images.build( - path=str(current_dir), - dockerfile="Dockerfile.postgres-hypopg", - buildargs={"PG_VERSION": pg_version, "PG_MAJOR": pg_version}, - tag=custom_image_name, - rm=True, - ) - logger.info(f"Successfully built image {custom_image_name}") - except Exception as e: - logger.error(f"Failed to build Docker image: {e}") - pytest.skip(f"Failed to build Docker image: {e}") - - postgres_password = "test_password" - postgres_db = "test_db" - - # Create container with more verbose logging - container = client.containers.run( - custom_image_name, - name=container_name, - environment={ - "POSTGRES_PASSWORD": postgres_password, - "POSTGRES_DB": postgres_db, - "POSTGRES_HOST_AUTH_METHOD": "trust", # Make authentication easier in tests - }, - ports={"5432/tcp": ("127.0.0.1", 0)}, # Let Docker assign a random port - command=[ - "-c", - "shared_preload_libraries=pg_stat_statements", - "-c", - "pg_stat_statements.track=all", - "-c", - "log_min_messages=info", # More verbose logging - "-c", - "log_statement=all", # Log all SQL statements - ], - detach=True, - ) - - logger.info(f"Container {container_name} started, waiting for PostgreSQL to be ready") - - try: - # Wait for container to start and get logs - time.sleep(2) # Give container a moment to start - container.reload() - - # Check if container is running - if container.status != "running": - logs = container.logs().decode("utf-8") - logger.error(f"Container {container_name} failed to start. Logs:\n{logs}") - pytest.skip(f"PostgreSQL container failed to start: {logs[:500]}...") - - # Get assigned port - port = container.ports["5432/tcp"][0]["HostPort"] - - # Wait for PostgreSQL to be ready - deadline = time.time() + 60 # Increased timeout to 60 seconds - is_ready = False - last_error = None - - while time.time() < deadline and not is_ready: - try: - exit_code, output = container.exec_run("pg_isready") - if exit_code == 0: - logger.info(f"PostgreSQL in container {container_name} is ready") - is_ready = True - break - else: - last_error = output.decode("utf-8") - logger.warning(f"PostgreSQL not ready yet: {last_error}") - except Exception as e: - last_error = str(e) - logger.warning(f"Error checking if PostgreSQL is ready: {e}") - - # Get container logs for debugging - if time.time() - deadline + 60 > 50: # Log when we're close to timeout - logs = container.logs().decode("utf-8") - logger.warning(f"Still waiting for PostgreSQL. Container logs:\n{logs[-2000:]}") - - time.sleep(2) - - if not is_ready: - logs = container.logs().decode("utf-8") - logger.error(f"Timeout waiting for PostgreSQL. Container logs:\n{logs[-2000:]}") - pytest.skip(f"Timeout waiting for PostgreSQL to start: {last_error}") - - connection_string = f"postgresql://postgres:{postgres_password}@localhost:{port}/{postgres_db}" - logger.info(f"PostgreSQL connection string: {connection_string}") - - yield connection_string, version - - except Exception as e: - logger.error(f"Error setting up PostgreSQL container: {e}") - # Get container logs for debugging - try: - logs = container.logs().decode("utf-8") - logger.error(f"Container logs:\n{logs}") - except Exception: - pass - raise - - finally: - logger.info(f"Stopping and removing container {container_name}") - try: - container.stop(timeout=1) - container.remove(v=True) - except Exception as e: - logger.warning(f"Error cleaning up container {container_name}: {e}") diff --git a/tools/generate_doc_tables.py b/tools/generate_doc_tables.py index 5b707d3..672e0ed 100755 --- a/tools/generate_doc_tables.py +++ b/tools/generate_doc_tables.py @@ -263,7 +263,6 @@ def render_tool_index() -> str: "mcpg.config": "Env-driven, validated `Settings` (frozen dataclass); redacts secrets in `__repr__`.", "mcpg.context": "`AppContext` — per-server state (settings, DB, cursor/listen managers) shared with tool wrappers.", "mcpg.schema_diff": "Structural schema diff powering `compare_schemas`.", - "mcpg._vendor": "Vendored MIT-licensed `SafeSqlDriver` + connection-pool kernel (SQL parse / allowlist / bind).", } @@ -285,10 +284,6 @@ def module_descriptions() -> dict[str, str]: parts = rel.parts if parts[-1] in ("__init__", "__main__"): continue - # Collapse the vendored SQL kernel into one aggregate row. - if parts[0] == "_vendor": - out.setdefault("mcpg._vendor", _MODULE_FALLBACK["mcpg._vendor"]) - continue name = "mcpg." + ".".join(parts) try: doc = ast.get_docstring(ast.parse(path.read_text(encoding="utf-8"))) or "" @@ -296,8 +291,6 @@ def module_descriptions() -> dict[str, str]: doc = "" desc = _first_sentence(doc) or _MODULE_FALLBACK.get(name, "") out[name] = _MODULE_FALLBACK.get(name, desc) if not desc else desc - # Ensure the aggregate vendor row exists even if walk order missed it. - out.setdefault("mcpg._vendor", _MODULE_FALLBACK["mcpg._vendor"]) return out