feat(subscriptions): add postgres connection and alembic migrations

.envrc

@@ -25,4 +25,7 @@ fi
 export VIRTUAL_ENV="${PWD}/.venv"
 PATH_add "${PWD}/.venv/bin"
 
+export SUBSCRIPTIONS_POSTGRES_URL="postgresql+psycopg://[email protected]:5432/snuba_subscriptions"
+export SUBSCRIPTIONS_POSTGRES_MIGRATIONS_URL="${SUBSCRIPTIONS_POSTGRES_URL}"
+
 . scripts/rust-envvars

.github/workflows/ci.yml

@@ -262,6 +262,7 @@ jobs:
             "ghcr.io/getsentry/image-mirror-confluentinc-cp-kafka:6.2.0"
             "ghcr.io/getsentry/image-mirror-confluentinc-cp-zookeeper:6.2.0"
             "ghcr.io/getsentry/docker-redis-cluster:7.0.10"
+            "ghcr.io/getsentry/image-mirror-library-postgres:18-alpine"
           )
           pulled=0
           for image in "${images[@]}"; do
@@ -334,6 +335,7 @@ jobs:
             "ghcr.io/getsentry/image-mirror-confluentinc-cp-kafka:6.2.0"
             "ghcr.io/getsentry/image-mirror-confluentinc-cp-zookeeper:6.2.0"
             "ghcr.io/getsentry/docker-redis-cluster:7.0.10"
+            "ghcr.io/getsentry/image-mirror-library-postgres:18-alpine"
           )
           for image in "${images[@]}"; do
             if ! docker image inspect "$image" &>/dev/null; then
@@ -528,6 +530,7 @@ jobs:
             "ghcr.io/getsentry/image-mirror-confluentinc-cp-kafka:6.2.0"
             "ghcr.io/getsentry/image-mirror-confluentinc-cp-zookeeper:6.2.0"
             "ghcr.io/getsentry/docker-redis-cluster:7.0.10"
+            "ghcr.io/getsentry/image-mirror-library-postgres:18-alpine"
           )
           for image in "${images[@]}"; do
             if ! docker image inspect "$image" &>/dev/null; then

devservices/config.yml

@@ -17,6 +17,12 @@ x-sentry-service-config:
         repo_name: sentry-shared-kafka
         branch: main
         repo_link: https://github.com/getsentry/sentry-shared-kafka.git
+    postgres:
+      description: Shared instance of postgres used by sentry services
+      remote:
+        repo_name: sentry-shared-postgres
+        branch: main
+        repo_link: https://github.com/getsentry/sentry-shared-postgres.git
     snuba:
       description: Service that provides a rich data model on top of Clickhouse together with a fast ingestion consumer and a query optimizer
     profiles-consumer:
@@ -28,7 +34,7 @@ x-sentry-service-config:
     metrics-consumer:
       description: Consumer for metrics
   modes:
-    default: [redis, kafka, clickhouse]
+    default: [redis, kafka, clickhouse, postgres]
     containerized: [clickhouse, redis, kafka, snuba]
     containerized-profiles: [clickhouse, redis, kafka, snuba, profiles-consumer, profile-chunks-consumer, functions-consumer, functions-consumer]
     containerized-metrics-dev: [clickhouse, redis, kafka, snuba, metrics-consumer]

docker-compose.gcb.yml

@@ -11,6 +11,8 @@ x-test-common: &test-common
       condition: service_started
     zookeeper:
       condition: service_started
+    postgres:
+      condition: service_healthy
   image: "$SNUBA_IMAGE"
   profiles: ["run_test"]
   volumes:
@@ -22,6 +24,7 @@ x-test-common: &test-common
     REDIS_HOST: "redis-cluster"
     REDIS_PORT: 7000
     REDIS_DB: 0
+    POSTGRES_HOST: postgres
     DEFAULT_BROKERS: "kafka:9092"
     SNUBA_TEST_SHARD: "${SNUBA_TEST_SHARD:-0}"
     SNUBA_TEST_SHARD_TOTAL: "${SNUBA_TEST_SHARD_TOTAL:-1}"
@@ -189,6 +192,16 @@ services:
         soft: 262144
         hard: 262144
 
+  postgres:
+    image: ghcr.io/getsentry/image-mirror-library-postgres:18-alpine
+    environment:
+      POSTGRES_HOST_AUTH_METHOD: trust
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U postgres"]
+      interval: 5s
+      timeout: 5s
+      retries: 3
+
   # This is the same thing as https://github.com/Grokzen/docker-redis-cluster but mirrored
   redis-cluster:
     image: ghcr.io/getsentry/docker-redis-cluster:7.0.10

pyproject.toml

@@ -5,6 +5,7 @@ version = "26.11.0.dev0"
 # does not resolve future Python versions that our pinned wheels do not support.
 requires-python = ">=3.13,<3.14"
 dependencies = [
+    "alembic>=1.20.0",
     "blinker>=1.9",
     "click>=8.1.7",
     "clickhouse-connect>=1.6.0",
@@ -26,6 +27,7 @@ dependencies = [
     "progressbar2>=4.2.0",
     "proto-plus>=1.24.0",
     "protobuf>=5.29.5",
+    "psycopg[binary]>=3.3.6",
     "pyasn1>=0.6.4",
     "python-dateutil>=2.8.2",
     "pyjwt[crypto]>=2.8.0",
@@ -44,6 +46,7 @@ dependencies = [
     "simplejson>=3.17.6",
     "snuba-sdk>=3.0.39",
     "sql-metadata>=3.0.1",
+    "sqlalchemy>=2.1.4",
     "sqlparse>=0.6.0",
     "structlog>=22.3.0",
     "structlog-sentry>=2.0.0",

snuba/cli/devserver.py

@@ -8,6 +8,7 @@
 import click
 
 from snuba import settings
+from snuba.subscriptions.postgres.database import create_database_if_missing
 
 # Match honcho: SIGTERM, then SIGKILL if children do not exit (avoids indefinite
 # hang when a child ignores SIGTERM; see also PEP 475 / wait() retry after signals).
@@ -52,6 +53,13 @@ def devserver(*, bootstrap: bool, workers: bool, log_level: str) -> None:
         if returncode > 0:
             sys.exit(returncode)
 
+        migrations_url = settings.SUBSCRIPTIONS_POSTGRES_MIGRATIONS_URL
+        if migrations_url is not None:
+            create_database_if_missing(migrations_url)
+            returncode = call(["snuba", "subscriptions-postgres", "migrate"])
+            if returncode > 0:
+                sys.exit(returncode)
+
     daemons = [("api", ["snuba", "api"])]
 
     if not workers:

snuba/cli/subscriptions_postgres.py

@@ -0,0 +1,14 @@
+import click
+
+from snuba.subscriptions.postgres.database import migrate_to_head
+
+
+@click.group()
+def subscriptions_postgres() -> None:
+    pass
+
+
+@subscriptions_postgres.command()
+def migrate() -> None:
+    """Upgrade the subscriptions Postgres database to the latest migration."""
+    migrate_to_head()

snuba/settings/__init__.py

@@ -181,6 +181,10 @@ class RedisClusterConfig(TypedDict):
 # default redis command timeout in seconds for redis commands (e.g. configs, rate limits) which are meant to be quick and fail-open
 REDIS_SOCKET_TIMEOUT = 0.1
 
+# The app connects as snuba through pgbouncer; migrations connect as snuba-migrations.
+SUBSCRIPTIONS_POSTGRES_URL = os.environ.get("SUBSCRIPTIONS_POSTGRES_URL")
+SUBSCRIPTIONS_POSTGRES_MIGRATIONS_URL = os.environ.get("SUBSCRIPTIONS_POSTGRES_MIGRATIONS_URL")
+
 
 class RedisClusters(TypedDict):
     cache: RedisClusterConfig | None

snuba/settings/settings_test.py

@@ -64,3 +64,9 @@
 VALIDATE_DATASET_YAMLS_ON_STARTUP = True
 
 LOG_MIGRATIONS = False
+
+POSTGRES_HOST = os.environ.get("POSTGRES_HOST", "127.0.0.1")
+SUBSCRIPTIONS_POSTGRES_URL = (
+    f"postgresql+psycopg://postgres@{POSTGRES_HOST}:5432/snuba_subscriptions_test"
+)
+SUBSCRIPTIONS_POSTGRES_MIGRATIONS_URL = SUBSCRIPTIONS_POSTGRES_URL

snuba/subscriptions/postgres/alembic.ini

@@ -0,0 +1,3 @@
+[alembic]
+script_location = %(here)s/migrations
+file_template = %%(rev)s_%%(slug)s

snuba/subscriptions/postgres/database.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+from typing import Any
+
+from alembic import command
+from alembic.config import Config
+from sqlalchemy import Engine, create_engine, make_url, text
+
+ALEMBIC_INI = Path(__file__).parent / "alembic.ini"
+
+CONNECT_ARGS: dict[str, Any] = {
+    # pgbouncer runs the app role in transaction mode, which breaks prepared statements.
+    "prepare_threshold": None,
+    "connect_timeout": 5,
+    # Detect a dead connection in under a minute instead of the OS default of 2 hours.
+    "keepalives_idle": 10,
+    "tcp_user_timeout": 30000,
+}
+
+
+def build_engine(url: str) -> Engine:
+    return create_engine(url, pool_pre_ping=True, connect_args=CONNECT_ARGS)
+
+
+def create_database_if_missing(url: str) -> None:
+    """Create the URL's database if absent; dev and test only, prod's comes from alloydb-init."""
+    target = make_url(url)
+    maintenance_url = target.set(database="postgres")
+    engine = create_engine(maintenance_url, isolation_level="AUTOCOMMIT")
+    query = text("SELECT 1 FROM pg_database WHERE datname = :name")
+    with engine.connect() as connection:
+        result = connection.execute(query, {"name": target.database})
+        exists = result.scalar()
+        if exists is None:
+            create = text(f'CREATE DATABASE "{target.database}"')
+            connection.execute(create)
+    engine.dispose()
+
+
+def migrate_to_head() -> None:
+    config = Config(str(ALEMBIC_INI))
+    command.upgrade(config, "head")

snuba/subscriptions/postgres/migrations/env.py

@@ -0,0 +1,18 @@
+from alembic import context
+
+from snuba import settings
+from snuba.subscriptions.postgres.database import build_engine
+
+if context.is_offline_mode():
+    raise RuntimeError("Offline migrations are not supported")
+
+url = settings.SUBSCRIPTIONS_POSTGRES_MIGRATIONS_URL
+if url is None:
+    raise RuntimeError("SUBSCRIPTIONS_POSTGRES_MIGRATIONS_URL is not set")
+
+engine = build_engine(url)
+with engine.connect() as connection:
+    context.configure(connection=connection, target_metadata=None)
+    with context.begin_transaction():
+        context.run_migrations()
+engine.dispose()

snuba/subscriptions/postgres/migrations/script.py.mako

@@ -0,0 +1,26 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+${imports if imports else ""}
+
+revision: str = ${repr(up_revision)}
+down_revision: str | Sequence[str] | None = ${repr(down_revision)}
+branch_labels: str | Sequence[str] | None = ${repr(branch_labels)}
+depends_on: str | Sequence[str] | None = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}

snuba/subscriptions/postgres/migrations/versions/0001_create_subscriptions.py

@@ -0,0 +1,38 @@
+"""Create the subscriptions table.
+
+Revision ID: 0001
+Revises:
+Create Date: 2026-10-16
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+revision: str = "0001"
+down_revision: str | Sequence[str] | None = None
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "subscriptions",
+        sa.Column("id", postgresql.UUID(), primary_key=True),
+        sa.Column("entity", sa.Text(), nullable=False),
+        sa.Column("partition", sa.Integer(), nullable=False),
+        sa.Column("data", postgresql.JSONB(), nullable=False),
+    )
+    op.create_index(
+        "subscriptions_entity_partition_idx",
+        "subscriptions",
+        ["entity", "partition"],
+    )
+
+
+def downgrade() -> None:
+    op.drop_index("subscriptions_entity_partition_idx", table_name="subscriptions")
+    op.drop_table("subscriptions")

tests/subscriptions/test_postgres.py

@@ -0,0 +1,49 @@
+from collections.abc import Generator
+
+import psycopg
+import pytest
+from sqlalchemy import Engine, inspect
+
+from snuba import settings
+from snuba.subscriptions.postgres.database import (
+    build_engine,
+    create_database_if_missing,
+    migrate_to_head,
+)
+
+
+@pytest.fixture(scope="module")
+def engine() -> Generator[Engine]:
+    url = settings.SUBSCRIPTIONS_POSTGRES_URL
+    assert url is not None
+    create_database_if_missing(url)
+    migrate_to_head()
+    engine = build_engine(url)
+    yield engine
+    engine.dispose()
+
+
+def test_subscriptions_table_schema(engine: Engine) -> None:
+    inspector = inspect(engine)
+    columns = inspector.get_columns("subscriptions")
+    shapes = {column["name"]: (str(column["type"]), column["nullable"]) for column in columns}
+    assert shapes == {
+        "id": ("UUID", False),
+        "entity": ("TEXT", False),
+        "partition": ("INTEGER", False),
+        "data": ("JSONB", False),
+    }
+
+    primary_key = inspector.get_pk_constraint("subscriptions")
+    assert primary_key["constrained_columns"] == ["id"]
+
+    indexes = inspector.get_indexes("subscriptions")
+    index_columns = {index["name"]: index["column_names"] for index in indexes}
+    assert index_columns == {"subscriptions_entity_partition_idx": ["entity", "partition"]}
+
+
+def test_engine_disables_prepared_statements(engine: Engine) -> None:
+    with engine.connect() as connection:
+        dbapi_connection = connection.connection.dbapi_connection
+        assert isinstance(dbapi_connection, psycopg.Connection)
+        assert dbapi_connection.prepare_threshold is None