From 6e6bca151810d093d28fd1a619216af7341135f4 Mon Sep 17 00:00:00 2001
From: momentum dev <dev@momentum.local>
Date: Fri, 17 Jul 2026 20:49:55 -0500
Subject: [PATCH] fix(attribution): stop KeyError-ing when DATABASE_URL is
 unset
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 research/attribution.py                 | 13 ++++---
 tests/unit/research/test_attribution.py | 52 +++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 5 deletions(-)
 create mode 100644 tests/unit/research/test_attribution.py

diff --git a/research/attribution.py b/research/attribution.py
index c5f7c1e..7669571 100644
--- a/research/attribution.py
+++ b/research/attribution.py
@@ -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:
diff --git a/tests/unit/research/test_attribution.py b/tests/unit/research/test_attribution.py
new file mode 100644
index 0000000..87f8765
--- /dev/null
+++ b/tests/unit/research/test_attribution.py
@@ -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