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:
1 parent
9772aca6bc
commit
6e6bca1518
2 files changed
+60
-5
No files matched your search
@@ -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:
|
||||
|
||||
@@ -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
|
||||
Reference in new issue
Block a user