Skip to content

feat: sql analysis - #1317

Open
JuanLeee wants to merge 4 commits into
devfrom
feat/kms-sql-analysis
Open

JuanLeee wants to merge 4 commits into
devfrom
feat/kms-sql-analysis

Conversation

@JuanLeee

@JuanLeee JuanLeee commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

Description

The KMS Encryption plugin has to know which bind parameter fills or is compared with which table.column. psycopg and MySQL Connector/Python both use pyformat placeholders, but they find them differently, and a placeholder attributed to the wrong column would encrypt the wrong value or let a plaintext one through. This adds that analysis as pure functions: a tokenizer that finds placeholders exactly as each driver does, and an analyzer that parses the statement with sqlglot. It is based on #1300 and extends its SqlDialect and MetadataManager. The plugin that uses it comes in #1331, so this PR has no plugin or interception code and changes nothing in the wrapper core.

  1. sql_tokenizer.py models each driver's placeholder rules. The drivers disagree on %%, on when they substitute at all and on which % sequences they reject, and neither skips quoted literals or comments, so a %s inside '...' is still a bound parameter. The tokenizer counts it, since skipping it would shift every later value onto the wrong column, and marks it hidden, which makes the analysis incomplete. For such a statement the analysis still lists the tables, if sqlglot can parse it, so the plugin can refuse the statement when one of them has encrypted columns. It also reads /*@encrypt:table.column*/ annotations, detects MySQL executable comments (/*! */, /*M! */), records charset introducers (_binary %s, which Django renders for binary values), and writes a normalized copy of the SQL with each placeholder renamed :_kms_p<ordinal>. It has no third-party dependencies.
  2. sql_analyzer.py parses the normalized copy and returns a frozen StatementAnalysis: the statement kind and whether it writes, the tables it names and writes, assignments (placeholders written into a column), comparisons (placeholders compared with a column), annotations, unbound_writes (columns written with a literal, an expression, DEFAULT, another column or a COPY stream), result_columns (the table column each select-list or RETURNING item comes from), and gap, which says why the analysis is incomplete. parameter_columns() gives the column for each ordinal, with annotations taking precedence. It covers INSERT (multi-row, INSERT ... SET, MySQL REPLACE), ON CONFLICT and ON DUPLICATE KEY UPDATE (row aliases included), UPDATE (PG UPDATE ... FROM, MySQL multi-table), DELETE, MERGE, data-modifying CTEs, COPY, PG PREPARE ... AS, several statements in one string, and SQLAlchemy's multi-row insert form (INSERT ... SELECT ... FROM (VALUES ...) AS v(...)).
  3. A select-list column read through a derived table or CTE is followed to the table column it comes from, up to 8 levels deep, so it still decrypts by position; SQLAlchemy's legacy Query.first() with a joined eager load reads its columns that way. A CTE sees only the CTEs written before it, as both engines resolve a WITH without RECURSIVE. Set operations, *, expressions and sources whose alias lists column names leave a column unknown, and RETURNING lists are not followed.
  4. SqlDialect.fold_identifier returns the name the server resolves an identifier to, and SqlDialect.metadata_key the form a name is compared in against encryption_metadata. PostgreSQL lower-cases unquoted names (ASCII only) and compares exactly. MySQL keeps names as written and compares them case-insensitively, as the uncached metadata query's WHERE does under MySQL's default collations. MetadataManager's cached lookups now compare the same way, so on MySQL with lower_case_table_names=0, two tables whose names differ only in case share encryption metadata. A name the driver returns as bytes (psycopg on a SQL_ASCII database, Connector/Python with use_unicode=False) is decoded as UTF-8, and anything else raises MetadataError, since a name that never matches would make an encrypted column look like an ordinary one.

How it fails safe

When the analysis can't establish what a writing statement writes, it reports a gap, and the plugin passes the statement through to the database trigger. A hidden placeholder is the exception: the plugin refuses the statement when it names a table with encrypted columns. When the analysis sees an encrypted column written with something other than a bind parameter, it reports an unbound write, and the plugin refuses the statement. The tests try hardest to rule out a placeholder attributed to the wrong column.

  1. A parse is trusted only if every placeholder the driver binds appears in it exactly once; otherwise the gap is PLACEHOLDER_MISMATCH. The ordinal travels in each placeholder's name because sqlglot's tree order isn't source order, and ? is a jsonb operator on PostgreSQL.
  2. sqlglot reads UPDATE LOW_PRIORITY users as a table named LOW_PRIORITY, so MySQL statement modifiers are removed before parsing, and a leading REPLACE is read as INSERT.
  3. Apart from the table listing for a hidden placeholder, only statements with known leading words are parsed. Anything else (EXPLAIN ANALYZE INSERT, EXECUTE, LOAD DATA) is classified from its keywords. sqlglot logs statement text, literals included, when it falls back to a generic command, so a thread-local filter drops its log records while this module parses, and detail never contains SQL text.
  4. In an upsert, copying the incoming row's own value (EXCLUDED.ssn, VALUES(ssn) or a MySQL row alias) is not an unbound write.
  5. A bare annotation (/*@encrypt:ssn*/) gets a table only when the statement says which: the table already inferred for that placeholder, else the single table a written value is written to, else the only table a compared value's query reads. A written value that reaches the write through a CTE or a FROM, JOIN, LATERAL or USING source stays bare. A value compared or tested inside a written expression (a CASE condition, the condition of IF(), a regex or JSON operator) counts as compared, so it never takes the written table. An annotation accepts every [\w.]+ reference that ends in a column, and quoted parts; one that ends in . names no column and is reported as stray.
  6. Results are cached (lru_cache(1024)) only for SQL of at most 4096 characters and 64 % signs. A RecursionError result is never cached, because it depends on the caller's stack depth.

Dependency

sqlglot is optional, like cryptography. It is installed only with the encryption extra (pip install "aws-advanced-python-wrapper[encryption]") and imported only by sql_analyzer.py, so the wrapper and the tokenizer import without it, and importing the analyzer without it raises an ImportError that names the extra. The pin is >=30.14.0,<31.0.0, with 30.25.0 in the lock file. It stops below the next major because the analysis walks sqlglot's expression tree, which a major release can change, so moving to a new major is a dependency bump that has to pass the analyzer's tests. The dev dependency group also has sqlglot, so the unit tests run without the extra.

Testing

test_sql_tokenizer.py (135 tests) uses the drivers' own placeholder scanners as an oracle, so a driver upgrade that changes binding fails there. test_sql_analyzer.py (373) covers the shapes above and the result-column tracing, including nesting to 8 levels and one past it and which CTEs a CTE sees, with and without RECURSIVE. test_metadata_manager.py gains 10 cases and test_encryption_import_guard.py 4. pytest ./tests/unit -Werror passes 3940 tests, and flake8, isort and mypy are clean on Python 3.10 and 3.12.

Through the plugin in #1331, the analysis also ran on Aurora PostgreSQL 17.9 and Aurora MySQL 3.10.3 with psycopg and both mysql-connector builds: every supported type, NULL, executemany, every fetch shape and loop, UPDATE, SQLAlchemy Core and ORM on both engines and Django on MySQL. Every stored value, read back without the plugin, was an encrypted payload.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

@JuanLeee

JuanLeee commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor Author

CI doesn't run on this PR while its base is #1300's branch, so here are the local results for 1276122 (the same tree as b82e57a, with the commit message reduced to its subject line), from a clean git archive export with poetry install -E encryption on Python 3.12.15:

  • flake8 .: clean
  • isort --check-only .: clean
  • mypy .: no issues in 389 source files
  • pytest ./tests/unit -Werror: 3448 passed
  • poetry check --lock: passes

@JuanLeee
JuanLeee force-pushed the feat/kms-sql-analysis branch from b82e57a to 1276122 Compare October 28, 2026 12:00
@JuanLeee
JuanLeee force-pushed the feat/encryption-service branch from a633a43 to 0f538c7 Compare October 28, 2026 12:00
@JuanLeee
JuanLeee force-pushed the feat/kms-sql-analysis branch from 1276122 to 4e4c9f9 Compare October 29, 2026 12:02
@JuanLeee
JuanLeee force-pushed the feat/encryption-service branch 2 times, most recently from 2296d47 to 9295923 Compare November 2, 2026 11:01
@JuanLeee
JuanLeee force-pushed the feat/kms-sql-analysis branch from 4e4c9f9 to 2cb1470 Compare November 2, 2026 11:01
@JuanLeee
JuanLeee force-pushed the feat/encryption-service branch from 9295923 to 89f9130 Compare November 5, 2026 00:02
@JuanLeee
JuanLeee force-pushed the feat/kms-sql-analysis branch from 2cb1470 to 448d68c Compare November 5, 2026 00:02
Base automatically changed from feat/encryption-service to feat/kms-stack-1-call-context November 12, 2026 09:38
@JuanLeee
JuanLeee force-pushed the feat/kms-stack-1-call-context branch from 765339a to a5e6595 Compare November 12, 2026 09:56
Base automatically changed from feat/kms-stack-1-call-context to dev November 12, 2026 10:04

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants