From b13470f390a27bb1e7a91d2c087d175aacc4ddaf Mon Sep 17 00:00:00 2001 From: jakeross Date: Thu, 15 Jan 2026 10:10:47 -0700 Subject: [PATCH 1/2] Unify DB initialization and app_read grants --- alembic/env.py | 4 ++-- db/initialization.py | 40 ++++++++++++++++++++++++++++++++++++++++ tests/conftest.py | 42 +++++++----------------------------------- transfers/transfer.py | 14 ++++++-------- 4 files changed, 55 insertions(+), 45 deletions(-) create mode 100644 db/initialization.py diff --git a/alembic/env.py b/alembic/env.py index f0bd9e778..5b7b03d43 100644 --- a/alembic/env.py +++ b/alembic/env.py @@ -5,9 +5,8 @@ from alembic import context from dotenv import load_dotenv -from sqlalchemy import create_engine, engine_from_config, pool, text - from services.util import get_bool_env +from sqlalchemy import create_engine, engine_from_config, pool, text # this is the Alembic Config object, which provides # access to the values within the .ini file in use. @@ -167,6 +166,7 @@ def getconn(): ).first() if not role_exists: autocommit_role.execute(text("CREATE ROLE app_read")) + autocommit_role.execute(text("GRANT app_read TO PUBLIC")) with connectable.connect() as connection: context.configure( diff --git a/db/initialization.py b/db/initialization.py new file mode 100644 index 000000000..d3847ce5e --- /dev/null +++ b/db/initialization.py @@ -0,0 +1,40 @@ +"""Utilities for preparing and keeping the database schema in sync.""" + +from db import Base +from sqlalchemy import text +from sqlalchemy.orm import Session +from sqlalchemy_searchable import sync_trigger +from sqlalchemy_utils import TSVectorType + +APP_READ_GRANT_SQL = text( + """ + DO $$ + BEGIN + IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'app_read') THEN + EXECUTE 'GRANT USAGE ON SCHEMA public TO app_read'; + EXECUTE 'GRANT SELECT ON ALL TABLES IN SCHEMA public TO app_read'; + EXECUTE 'ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO app_read'; + EXECUTE 'GRANT app_read TO PUBLIC'; + END IF; + END $$; + """ +) + + +def recreate_public_schema(session: Session) -> None: + """Drop and recreate the public schema, PostGIS extension, and app_read grants.""" + session.execute(text("DROP SCHEMA public CASCADE")) + session.execute(text("CREATE SCHEMA public")) + session.execute(text("CREATE EXTENSION IF NOT EXISTS postgis")) + session.execute(APP_READ_GRANT_SQL) + session.commit() + + +def sync_search_vector_triggers(session: Session) -> None: + """Ensure SQLAlchemy-searchable triggers exist for every TSVector column.""" + conn = session.connection() + for table in Base.metadata.tables.values(): + for column in table.columns: + if isinstance(column.type, TSVectorType): + sync_trigger(conn, table.name, column.name, list(column.type.columns)) + session.commit() diff --git a/tests/conftest.py b/tests/conftest.py index 979778cb3..dbec0fa2d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,17 +1,16 @@ import os -import uuid import pytest from alembic import command from alembic.config import Config -from dotenv import load_dotenv -from sqlalchemy import text -from sqlalchemy_utils import TSVectorType -from sqlalchemy_searchable import sync_trigger - from core.initializers import init_lexicon, init_parameter from db import * from db.engine import session_ctx +from db.initialization import ( + recreate_public_schema, + sync_search_vector_triggers, +) +from dotenv import load_dotenv from tests import get_parameter_id @@ -29,39 +28,12 @@ def _alembic_config() -> Config: def _reset_schema() -> None: with session_ctx() as session: - session.execute(text("DROP SCHEMA public CASCADE")) - session.execute(text("CREATE SCHEMA public")) - session.execute(text("CREATE EXTENSION IF NOT EXISTS postgis")) - session.execute( - text( - """ - DO $$ - BEGIN - IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'app_read') THEN - EXECUTE 'GRANT USAGE ON SCHEMA public TO app_read'; - EXECUTE 'GRANT SELECT ON ALL TABLES IN SCHEMA public TO app_read'; - EXECUTE 'ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO app_read'; - END IF; - END $$; - """ - ) - ) - session.commit() + recreate_public_schema(session) def _sync_search_vectors() -> None: with session_ctx() as session: - conn = session.connection() - for table in Base.metadata.tables.values(): - for column in table.columns: - if isinstance(column.type, TSVectorType): - sync_trigger( - conn, - table.name, - column.name, - list(column.type.columns), - ) - session.commit() + sync_search_vector_triggers(session) @pytest.fixture(scope="session", autouse=True) diff --git a/transfers/transfer.py b/transfers/transfer.py index 0586e62bc..0d497efc0 100644 --- a/transfers/transfer.py +++ b/transfers/transfer.py @@ -17,13 +17,11 @@ import time from concurrent.futures import ThreadPoolExecutor, as_completed -from dotenv import load_dotenv - from alembic import command from alembic.config import Config - from db.engine import session_ctx -from sqlalchemy import text +from db.initialization import recreate_public_schema, sync_search_vector_triggers +from dotenv import load_dotenv from services.util import get_bool_env from transfers.aquifer_system_transfer import transfer_aquifer_systems from transfers.geologic_formation_transfer import transfer_geologic_formations @@ -145,10 +143,7 @@ def _alembic_config() -> Config: def _drop_and_rebuild_db() -> None: logger.info("Dropping schema public") with session_ctx() as session: - session.execute(text("DROP SCHEMA public CASCADE")) - session.execute(text("CREATE SCHEMA public")) - session.execute(text("CREATE EXTENSION IF NOT EXISTS postgis")) - session.commit() + recreate_public_schema(session) logger.info("Running Alembic migrations") try: command.upgrade(_alembic_config(), "head") @@ -159,6 +154,9 @@ def _drop_and_rebuild_db() -> None: "Alembic upgrade returned SystemExit(%s); continuing transfer", exc.code ) logger.info("Alembic migrations complete") + logger.info("Synchronizing search vector triggers") + with session_ctx() as session: + sync_search_vector_triggers(session) logger.info("Initializing lexicon data") init_lexicon() logger.info("Initializing parameter data") From 9063ac345fb06cc35b36fd2801b5e0f29411693c Mon Sep 17 00:00:00 2001 From: jakeross Date: Thu, 15 Jan 2026 10:20:34 -0700 Subject: [PATCH 2/2] Restrict app_read grants and share initialization --- alembic/env.py | 3 ++- db/initialization.py | 40 ++++++++++++++++++++++++++++++++++++++-- tests/conftest.py | 3 ++- transfers/transfer.py | 3 ++- 4 files changed, 44 insertions(+), 5 deletions(-) diff --git a/alembic/env.py b/alembic/env.py index 5b7b03d43..089144e88 100644 --- a/alembic/env.py +++ b/alembic/env.py @@ -24,6 +24,7 @@ # from db import Base # Import your Base from models/__init__.py from db import Base +from db.initialization import grant_app_read_members # target_metadata = mymodel.Base.metadata target_metadata = Base.metadata @@ -166,7 +167,7 @@ def getconn(): ).first() if not role_exists: autocommit_role.execute(text("CREATE ROLE app_read")) - autocommit_role.execute(text("GRANT app_read TO PUBLIC")) + grant_app_read_members(autocommit_role) with connectable.connect() as connection: context.configure( diff --git a/db/initialization.py b/db/initialization.py index d3847ce5e..862b61749 100644 --- a/db/initialization.py +++ b/db/initialization.py @@ -1,11 +1,15 @@ """Utilities for preparing and keeping the database schema in sync.""" -from db import Base +import os + from sqlalchemy import text +from sqlalchemy.engine import Connection from sqlalchemy.orm import Session from sqlalchemy_searchable import sync_trigger from sqlalchemy_utils import TSVectorType +from db import Base + APP_READ_GRANT_SQL = text( """ DO $$ @@ -14,12 +18,43 @@ EXECUTE 'GRANT USAGE ON SCHEMA public TO app_read'; EXECUTE 'GRANT SELECT ON ALL TABLES IN SCHEMA public TO app_read'; EXECUTE 'ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO app_read'; - EXECUTE 'GRANT app_read TO PUBLIC'; END IF; END $$; """ ) +GRANT_MEMBER_SQL = text( + """ + DO $$ + DECLARE + username text := :grantee; + BEGIN + IF username IS NULL OR username = '' THEN + RETURN; + END IF; + IF EXISTS (SELECT 1 FROM pg_roles WHERE rolname = username) THEN + EXECUTE format('GRANT app_read TO %I', username); + END IF; + END $$; + """ +) + + +def _parse_app_read_members() -> list[str]: + members = os.environ.get("APP_READ_MEMBERS", "") + return [member.strip() for member in members.split(",") if member.strip()] + + +def grant_app_read_members(executor: Session | Connection | None) -> None: + """Grant app_read to each configured role if it exists.""" + if executor is None: + return + members = _parse_app_read_members() + if not members: + return + for member in members: + executor.execute(GRANT_MEMBER_SQL, {"grantee": member}) + def recreate_public_schema(session: Session) -> None: """Drop and recreate the public schema, PostGIS extension, and app_read grants.""" @@ -27,6 +62,7 @@ def recreate_public_schema(session: Session) -> None: session.execute(text("CREATE SCHEMA public")) session.execute(text("CREATE EXTENSION IF NOT EXISTS postgis")) session.execute(APP_READ_GRANT_SQL) + grant_app_read_members(session) session.commit() diff --git a/tests/conftest.py b/tests/conftest.py index dbec0fa2d..032732740 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,7 @@ import os +from dotenv import load_dotenv + import pytest from alembic import command from alembic.config import Config @@ -10,7 +12,6 @@ recreate_public_schema, sync_search_vector_triggers, ) -from dotenv import load_dotenv from tests import get_parameter_id diff --git a/transfers/transfer.py b/transfers/transfer.py index 0d497efc0..687166733 100644 --- a/transfers/transfer.py +++ b/transfers/transfer.py @@ -17,11 +17,12 @@ import time from concurrent.futures import ThreadPoolExecutor, as_completed +from dotenv import load_dotenv + from alembic import command from alembic.config import Config from db.engine import session_ctx from db.initialization import recreate_public_schema, sync_search_vector_triggers -from dotenv import load_dotenv from services.util import get_bool_env from transfers.aquifer_system_transfer import transfer_aquifer_systems from transfers.geologic_formation_transfer import transfer_geologic_formations