fix(attribution): stop KeyError-ing when DATABASE_URL is unset

research/attribution.py's _connect() read os.environ["DATABASE_URL"]
directly, bypassing the storage facade fallback chain (DATABASE_URL,
PG* vars, LAN host retries) every other module goes through — and
KeyError-ing (uncaught by name, though still 500'd by the endpoint's
broad except) whenever DATABASE_URL was unset.

Delegate to storage.db.get_connection().raw instead, so DB resolution
matches EventStore/research.multi_schema, and a genuinely unreachable
DB fails with a catchable psycopg2.OperationalError rather than a
KeyError.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
momentum devandClaude Fable 5 committed 2026-07-17 20:49:55 -05:00
1 parent 9772aca6bc
commit 6e6bca1518
2 files changed
+60 -5

No files matched your search

+8 -5
View File
@@ -19,9 +19,8 @@ All queries filter on the native `timestamp` column first (payload_json casts on
import argparse
import os
from datetime import date, timedelta
from urllib.parse import urlparse
import psycopg2
from storage.db import get_connection
try: # optional; .env for CLI use
from dotenv import load_dotenv
@@ -33,9 +32,13 @@ CONFIRM_MINUTE_DEFAULT = 15 # mirror TRADING_ENTRY_CONFIRM_BY_MINUTE for the e
def _connect():
u = urlparse(os.environ["DATABASE_URL"])
return psycopg2.connect(host=u.hostname, port=u.port, user=u.username,
password=u.password, dbname=u.path.lstrip("/"))
"""Resolve the Postgres connection through the same facade every other
module uses (storage.db_pg's DATABASE_URL / PG* env fallback chain),
instead of hand-rolling `os.environ["DATABASE_URL"]` — which KeyErrors
(and 500s /api/attribution) whenever DATABASE_URL is unset. A genuinely
unreachable DB still raises (psycopg2.OperationalError), which the
endpoint's existing `except Exception` turns into a clean JSON error."""
return get_connection().raw
def compute_attribution(days: int = 14, confirm_minute: int | None = None) -> dict:
+52
View File
@@ -0,0 +1,52 @@
"""research/attribution.py _connect() must resolve the DB connection through
the same storage facade (storage.db_pg) every other module uses, instead of
reading os.environ["DATABASE_URL"] directly — which KeyErrors (and 500s the
/api/attribution endpoint) whenever DATABASE_URL is unset.
This test never touches the production database: it points PGHOST at a
nonexistent socket path so a genuinely unreachable DB fails fast and cleanly,
without a KeyError, and without any real connection attempt succeeding."""
from __future__ import annotations
import psycopg2
import pytest
from research import attribution
def test_connect_without_database_url_does_not_keyerror(monkeypatch, tmp_path):
monkeypatch.delenv("DATABASE_URL", raising=False)
# Redirect the no-DSN fallback path at a socket dir that cannot exist, so
# this test can never reach a real (production or otherwise) Postgres.
monkeypatch.setenv("PGHOST", str(tmp_path / "no-such-pg-socket-dir"))
monkeypatch.setenv("PG_SKIP_SCHEMA", "1")
with pytest.raises(Exception) as excinfo:
attribution._connect()
assert not isinstance(excinfo.value, KeyError)
# a clean, catchable connection failure — exactly what the endpoint's
# existing `except Exception` turns into a JSON error response.
assert isinstance(excinfo.value, psycopg2.OperationalError)
def test_connect_uses_storage_db_facade(monkeypatch):
"""The fix should delegate to storage.db.get_connection (the fallback
chain every other module uses) rather than hand-rolling env parsing."""
calls = []
sentinel = object()
def fake_get_connection(*a, **kw):
calls.append((a, kw))
class _Fake:
raw = sentinel
return _Fake()
monkeypatch.setattr(attribution, "get_connection", fake_get_connection)
conn = attribution._connect()
assert calls, "attribution._connect() must call storage.db.get_connection()"
assert conn is sentinel