add alembic instrumentation

Author: martinabeleda

alembic/env.py

@@ -1,6 +1,7 @@
 import os
 from logging.config import fileConfig
 
+import logfire
 from sqlalchemy import engine_from_config, pool
 
 from alembic import context
@@ -16,7 +17,9 @@
 
 # add your model's MetaData object here
 # for 'autogenerate' support
+from strava.config import settings  # noqa: E402
 from strava.db.base import Base  # noqa: E402
+from strava.observability import configure_migration_observability  # noqa: E402
 
 target_metadata = Base.metadata
 
@@ -66,6 +69,7 @@ def run_migrations_online():
         prefix="sqlalchemy.",
         poolclass=pool.NullPool,
     )
+    configure_migration_observability(settings, connectable)
 
     with connectable.connect() as connection:
         context.configure(
@@ -75,7 +79,8 @@ def run_migrations_online():
         )
 
         with context.begin_transaction():
-            context.run_migrations()
+            with logfire.span("alembic run migrations", db_url=configuration["sqlalchemy.url"]):
+                context.run_migrations()
 
 
 if context.is_offline_mode():

docker-compose.test.yml

@@ -18,6 +18,8 @@ services:
     build: ./
     command: [uv, run, alembic, upgrade, head]
     environment:
+      - PROJECT_NAME=strava-test
+      - ENVIRONMENT=test
       - POSTGRES_SERVER=db:5432
       - POSTGRES_USER=postgres
       - POSTGRES_PASSWORD=password

docker-compose.yml

@@ -22,6 +22,9 @@ services:
       - ./:/root/app
     working_dir: /root/app
     environment:
+      - PROJECT_NAME=strava
+      - ENVIRONMENT=development
+      - LOGFIRE_TOKEN=${LOGFIRE_TOKEN}
       - POSTGRES_SERVER=db:5432
       - POSTGRES_USER=postgres
       - POSTGRES_PASSWORD=password

pyproject.toml

@@ -19,7 +19,7 @@ dependencies = [
   "geoalchemy2",
   "alembic",
   "psycopg2",
-  "logfire>=4.29.0",
+  "logfire[system-metrics]>=4.29.0",
   "opentelemetry-instrumentation-fastapi>=0.60b1",
   "opentelemetry-instrumentation-sqlalchemy>=0.60b1",
 ]

strava/observability.py

@@ -5,12 +5,22 @@
 from strava.config import Settings
 
 
-def configure_observability(app: FastAPI, settings: Settings, engine: Engine) -> None:
+def configure_logfire(settings: Settings, *, service_name: str | None = None) -> None:
     logfire.configure(
-        service_name=settings.PROJECT_NAME,
+        service_name=service_name or settings.PROJECT_NAME,
         environment=settings.ENVIRONMENT,
         send_to_logfire=settings.LOGFIRE_SEND_TO_LOGFIRE,
         token=settings.LOGFIRE_TOKEN,
     )
+
+
+def configure_observability(app: FastAPI, settings: Settings, engine: Engine) -> None:
+    configure_logfire(settings)
     logfire.instrument_fastapi(app)
     logfire.instrument_sqlalchemy(engine)
+    logfire.instrument_system_metrics()
+
+
+def configure_migration_observability(settings: Settings, engine: Engine) -> None:
+    configure_logfire(settings, service_name=f"{settings.PROJECT_NAME}-alembic")
+    logfire.instrument_sqlalchemy(engine)

tests/test_config.py

@@ -27,10 +27,17 @@ def test_environment_default(self):
 
         assert settings.ENVIRONMENT == "development"
 
-    def test_logfire_token_default(self):
-        from strava.config import settings
+    def test_logfire_token_default(self, monkeypatch):
+        monkeypatch.setenv("PROJECT_NAME", "test")
+        monkeypatch.setenv("POSTGRES_SERVER", "localhost")
+        monkeypatch.setenv("POSTGRES_PASSWORD", "pass")
+        monkeypatch.setenv("POSTGRES_DB", "db")
+        monkeypatch.delenv("LOGFIRE_TOKEN", raising=False)
+
+        from strava.config import Settings
 
-        assert settings.LOGFIRE_TOKEN is None
+        s = Settings()
+        assert s.LOGFIRE_TOKEN is None
 
     def test_settings_from_env(self, monkeypatch):
         monkeypatch.setenv("PROJECT_NAME", "my-project")

tests/test_observability.py

@@ -3,21 +3,45 @@
 from fastapi import FastAPI
 
 from strava.config import Settings
-from strava.observability import configure_observability
+from strava.observability import (
+    configure_logfire,
+    configure_migration_observability,
+    configure_observability,
+)
+
+
+def build_settings(**overrides):
+    values = {
+        "PROJECT_NAME": "strava-test",
+        "ENVIRONMENT": "test",
+        "POSTGRES_SERVER": "localhost",
+        "POSTGRES_PASSWORD": "password",
+        "POSTGRES_DB": "test",
+        "LOGFIRE_SEND_TO_LOGFIRE": False,
+        "LOGFIRE_TOKEN": "test-token",
+    }
+    values.update(overrides)
+    return Settings(**values)
+
+
+def test_configure_logfire_uses_default_service_name():
+    settings = build_settings()
+
+    with patch("strava.observability.logfire.configure") as configure:
+        configure_logfire(settings)
+
+    configure.assert_called_once_with(
+        service_name="strava-test",
+        environment="test",
+        send_to_logfire=False,
+        token="test-token",
+    )
 
 
 def test_configure_observability_wires_logfire():
     app = FastAPI()
     engine = MagicMock()
-    settings = Settings(
-        PROJECT_NAME="strava-test",
-        ENVIRONMENT="test",
-        POSTGRES_SERVER="localhost",
-        POSTGRES_PASSWORD="password",
-        POSTGRES_DB="test",
-        LOGFIRE_SEND_TO_LOGFIRE=False,
-        LOGFIRE_TOKEN="test-token",
-    )
+    settings = build_settings()
 
     with patch("strava.observability.logfire.configure") as configure:
         with patch("strava.observability.logfire.instrument_fastapi") as instrument_fastapi:
@@ -34,3 +58,20 @@ def test_configure_observability_wires_logfire():
     )
     instrument_fastapi.assert_called_once_with(app)
     instrument_sqlalchemy.assert_called_once_with(engine)
+
+
+def test_configure_migration_observability_uses_alembic_service_name():
+    engine = MagicMock()
+    settings = build_settings()
+
+    with patch("strava.observability.logfire.configure") as configure:
+        with patch("strava.observability.logfire.instrument_sqlalchemy") as instrument_sqlalchemy:
+            configure_migration_observability(settings, engine)
+
+    configure.assert_called_once_with(
+        service_name="strava-test-alembic",
+        environment="test",
+        send_to_logfire=False,
+        token="test-token",
+    )
+    instrument_sqlalchemy.assert_called_once_with(engine)