feat(storage): SQLAlchemy storage provider with PostgreSQL support - #1161
feat(storage): SQLAlchemy storage provider with PostgreSQL support#1161vringar wants to merge 16 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
Adds a new SQLAlchemy-based structured storage provider and schema definitions to enable multi-backend SQL storage (SQLite + PostgreSQL), updates test scaffolding to include SQLAlchemy/PG scenarios, and adjusts generated test values to fit PostgreSQL INTEGER bounds for certain columns.
Changes:
- Introduces
SQLAlchemyStorageProvider(SQLAlchemy Core) and a shared SQLAlchemy schema (TABLE_MAP) for OpenWPM’s structured tables. - Updates storage test fixtures and adds SQLAlchemy provider tests, including optional PostgreSQL coverage gated on
pytest-postgresql. - Pins/adds PostgreSQL + psycopg2 + SQLAlchemy + pytest-postgresql dependencies; adjusts random test value ranges for
duration/response_status.
Reviewed changes
Copilot reviewed 9 out of 10 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
openwpm/storage/sqlalchemy_provider.py |
New SQLAlchemy-backed StructuredStorageProvider implementation (incl. reflection fallback + coercions). |
openwpm/storage/sqlalchemy_schema.py |
SQLAlchemy Core table definitions for the structured schema, with BigInteger choices for PG overflow avoidance. |
openwpm/storage/sql_provider.py |
Rewrites SQLiteStorageProvider into a thin wrapper delegating to SQLAlchemyStorageProvider. |
test/storage/test_sqlalchemy_provider.py |
New tests for schema equivalence, all-tables insert smoke tests, and _coerce_record behavior; adds optional PG scenario test. |
test/storage/fixtures.py |
Adds sqlalchemy_sqlite scenario and gated postgresql scenario to structured provider fixtures. |
test/storage/conftest.py |
Adds conditional pytest-postgresql fixtures for discovery when the dependency is installed. |
test/storage/test_values.py |
Narrows random ranges for duration/response_status to 32-bit int bounds. |
environment.yaml |
Adds pinned deps: PostgreSQL, psycopg2, pytest-postgresql, SQLAlchemy. |
scripts/environment-unpinned.yaml |
Adds corresponding unpinned dependency entries. |
.gitignore |
Ignores Crosslink-managed local state files. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| self._connection.execute(sa_table.insert(), record) | ||
| self._sql_counter += 1 | ||
| except Exception as e: | ||
| self.logger.error( | ||
| "Unsupported record:\n%s\n%s\ntable=%s\n%s\n" | ||
| % (type(e), e, table, repr(record)) | ||
| ) | ||
| # On PostgreSQL, a failed statement aborts the entire transaction. | ||
| # All subsequent statements would fail with "InFailedSqlTransaction" | ||
| # until a ROLLBACK is issued. We must rollback here so that | ||
| # subsequent inserts can succeed. | ||
| try: | ||
| self._connection.rollback() | ||
| except Exception as rollback_err: | ||
| self.logger.error("Rollback failed: %s", rollback_err) |
There was a problem hiding this comment.
store_record rolls back the entire transaction on any insert error. That rollback will also undo any previous successful inserts since the last commit, causing silent data loss for a whole batch/visit when a single bad record arrives. Consider isolating each insert with a SAVEPOINT / nested transaction (e.g., begin_nested()), or committing more frequently and only rolling back the failed statement, while still recovering from PostgreSQL's aborted-transaction state.
| self._connection.execute(sa_table.insert(), record) | |
| self._sql_counter += 1 | |
| except Exception as e: | |
| self.logger.error( | |
| "Unsupported record:\n%s\n%s\ntable=%s\n%s\n" | |
| % (type(e), e, table, repr(record)) | |
| ) | |
| # On PostgreSQL, a failed statement aborts the entire transaction. | |
| # All subsequent statements would fail with "InFailedSqlTransaction" | |
| # until a ROLLBACK is issued. We must rollback here so that | |
| # subsequent inserts can succeed. | |
| try: | |
| self._connection.rollback() | |
| except Exception as rollback_err: | |
| self.logger.error("Rollback failed: %s", rollback_err) | |
| # Isolate each insert in a SAVEPOINT so a single bad record does | |
| # not roll back earlier successful inserts in the outer transaction. | |
| # This also recovers cleanly from PostgreSQL's aborted-transaction | |
| # state by rolling back only the failed statement scope. | |
| with self._connection.begin_nested(): | |
| self._connection.execute(sa_table.insert(), record) | |
| self._sql_counter += 1 | |
| except Exception as e: | |
| self.logger.error( | |
| "Unsupported record:\n%s\n%s\ntable=%s\n%s\n" | |
| % (type(e), e, table, repr(record)) | |
| ) |
75e9c42 to
267ac02
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #1161 +/- ##
==========================================
+ Coverage 62.36% 62.79% +0.42%
==========================================
Files 40 42 +2
Lines 3930 3986 +56
==========================================
+ Hits 2451 2503 +52
- Misses 1479 1483 +4 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
46e2d22 to
cff8196
Compare
cff8196 to
6d09f0c
Compare
6d09f0c to
cf876d1
Compare
c1052c1 to
9fcbaeb
Compare
9fcbaeb to
a8bb697
Compare
Add SQLAlchemyStorageProvider backed by SQLAlchemy Core (not ORM) that supports any SQLAlchemy-compatible database. SQLiteStorageProvider is now a thin wrapper delegating to SQLAlchemyStorageProvider with a sqlite:/// URL. New files: - sqlalchemy_schema.py: All 13 tables as SQLAlchemy Table objects - sqlalchemy_provider.py: StructuredStorageProvider implementation - test_sqlalchemy_provider.py: Schema equivalence, all-tables, coercion tests Key decisions: - DATETIME columns → Text for cross-dialect compatibility - No foreign keys (test data violates them, SQLite ignores them) - PostgreSQL transaction abort recovery via rollback in store_record - Table reflection fallback for custom tables (e.g. page_links) Closes #1143
Add psycopg2, pytest-postgresql, and postgresql to environment.yaml and extend test/storage/fixtures.py with a postgresql scenario backed by SQLAlchemyStorageProvider, so the existing parametrized all-tables insertion tests also run against a real PostgreSQL instance.
…tures - Change Integer to BigInteger for columns holding large values - Gate PostgreSQL scenario so existing SQLite tests run without pg - Add missing deps to environment-unpinned.yaml
- Gate postgresql_scenarios on HAS_PYTEST_POSTGRESQL so SQLite tests work without pytest-postgresql installed - Remove psycopg2 from deps (mutually exclusive with psycopg2-binary) - Add explanatory comment for intentional ImportError pass in conftest
psycopg2 is a C extension that requires libpq, which is incompatible with python_abi 3.14 cp314t on conda-forge. psycopg (v3) is a pure Python driver that works on all Python versions. - Replace psycopg2 with psycopg in environment files - Remove postgresql server package (pytest-postgresql manages its own) - Update SQLAlchemy dialect from postgresql+psycopg2 to postgresql+psycopg
SQLite requires exactly INTEGER (not BIGINT) for AUTOINCREMENT primary keys. The SQLAlchemy schema used BigInteger for task_id, which rendered as BIGINT in DDL — causing OperationalError during table creation. Since SQLiteStorageProvider now delegates to SQLAlchemyStorageProvider, this crashed the StorageController process, and TaskManager blocked forever on status_queue.get(), hanging ALL test groups in CI. Changes: - task.task_id: BigInteger → Integer (matches schema.sql) - crawl.task_id: BigInteger → Integer (foreign key consistency) - test_values: cap task_id range to 2^31-1 for PostgreSQL compat - dns_responses: add missing redirect_url column from schema.sql
The SQLAlchemy dns_responses table was missing the error column present in schema.sql (added by the v0.35.0 DNS error-handling work). This left the SQLAlchemy schema at 11 columns vs schema.sql's 12, failing test_schema_equivalence and causing DNS-error data to silently not persist via the SQLAlchemy provider.
…storage provider Bump the SQLAlchemy storage provider's runtime dependencies (sqlalchemy, psycopg[binary]) to the versions verified working against a live PostgreSQL server, and classify the test-only pytest-postgresql dependency as a dev dependency. A fresh `conda env create -f environment.yaml` now yields a working `import sqlalchemy` / `import psycopg`.
pytest-postgresql is now a pinned (dev) dependency, so the package import succeeds wherever the environment is installed - including CI images that have no PostgreSQL server. Spinning up a real instance requires the server binaries on PATH (located via pg_config); without them the scenario errored instead of skipping. Gate the scenario on pg_config availability so the suite skips the PostgreSQL case cleanly when no server is present.
The Docker crawler entrypoint (crawler.py) hardcoded GCS for structured storage. Allow a real crawl to write structured data to PostgreSQL (or any SQLAlchemy-supported database) by setting OPENWPM_POSTGRES_URL (DB_URL accepted as an alias) to a SQLAlchemy DSN, e.g. postgresql+psycopg://user:pass@host:port/db. When set, the crawler uses SQLAlchemyStorageProvider for structured storage instead of GCS, mirroring the existing provider-selection pattern. demo.py documents the DSN form.
… provider Mirror the Arrow/parquet provider's per-provider-instance shard key so that many volunteer nodes can write to ONE shared (PostgreSQL) database without primary-key collisions or cross-client lock contention. visit_id and browser_id are per-node sequential, so two nodes both produce visit_id=1 -> on a shared database that is a primary-key collision and the second insert is lost. The SQLAlchemy provider now generates a per-instance instance_id (random 32-bit, matching arrow_storage) and stamps it onto every record before insert. - sqlalchemy_schema.py: add an instance_id column to every table. For tables keyed on a client-supplied id (crawl/browser_id, site_visits/visit_id, task/task_id, javascript/id, javascript_cookies/id) the primary key becomes composite with instance_id so cross-client inserts no longer collide. Tables with a database-generated autoincrement id (http_*, callstacks, dns_responses) and the keyless tables (crawl_history, navigations, incomplete_visits) just gain the column. - sqlalchemy_provider.py: generate self._instance_id per provider instance and stamp record[instance_id] before insert (guarded on column presence so reflected custom tables are unaffected); also stamp the direct incomplete_visits insert. - test_schema_equivalence: instance_id is a sharding column absent from schema.sql (the single-node SQLite path), so exclude it from the column comparison and compare primary-key membership instead of ordinal; task is exempt from the AUTOINCREMENT check (composite key cannot autoincrement and task_id is always client-supplied).
task and crawl have no visit_id column, so the test fixture's INVALID_VISIT_ID sentinel must be removed before the record reaches the provider (mirroring StorageController). Otherwise the unknown column makes the insert fail silently, and because store_record swallows exceptions the test would pass without storing those rows. Also correct the schema-equivalence test docstring, which claimed to compare column default values that it never actually checks.
store_record accumulates records in a single open transaction that is only committed on flush_cache/finalize_visit_id. On an insert failure it called self._connection.rollback(), which discards EVERY record buffered since the last commit (all visits for ~30s), not just the offending row — silent partial data loss on the default SQLite path, worse on PostgreSQL. Wrap each insert in a SAVEPOINT (begin_nested) so only the failing statement rolls back while the outer transaction (the good rows) survives. On PostgreSQL this also clears the aborted-transaction state so later inserts still succeed. SAVEPOINT is supported by both SQLite and PostgreSQL. Add read-back tests: one asserting stored values actually persist, one asserting a mid-batch failure (duplicate PK) drops only the bad row.
a8bb697 to
29cefcb
Compare
Summary
HAS_PYTEST_POSTGRESQLflagSupersedes #1149. Incorporates fixes from 5 rounds of adversarial review (VDD methodology).
VDD Review History
Test plan
pre-commit run --all-filespasses