test_data/test_db.py

"""
Standalone DB connectivity test — run with:
    uv run python test_db.py

Tests each layer independently so the exact failure point is visible.
Loads DATABASE_URL from .env automatically.
"""
import asyncio
import os
import socket
import sys
from urllib.parse import urlparse

from dotenv import load_dotenv

load_dotenv(override=True)

# asyncpg requires SelectorEventLoop on Windows — ProactorEventLoop (default in
# Python 3.8+ on Windows) causes asyncpg to hang silently on connect.
if sys.platform == "win32":
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

RAW_URL = os.environ.get("DATABASE_URL", "")

# ── 0. Show what URL we are working with ──────────────────────────────────────
print("\n── Step 0: URL check ──────────────────────────────────────────────────")
if not RAW_URL:
    print("FAIL  DATABASE_URL is not set in .env")
    sys.exit(1)

parsed = urlparse(RAW_URL)
masked = RAW_URL.replace(parsed.password or "", "***") if parsed.password else RAW_URL
print(f"  Raw scheme : {parsed.scheme}")
print(f"  Host       : {parsed.hostname}")
print(f"  Port       : {parsed.port or 5432}")
print(f"  Database   : {parsed.path.lstrip('/')}")
print(f"  Full URL   : {masked}")

# asyncpg needs postgresql+asyncpg:// — detect and auto-fix for this test
ASYNCPG_URL = RAW_URL
if parsed.scheme == "postgresql" or parsed.scheme == "postgres":
    ASYNCPG_URL = RAW_URL.replace(parsed.scheme + "://", "postgresql+asyncpg://", 1)
    print(f"\n  NOTE: scheme was '{parsed.scheme}://' — using 'postgresql+asyncpg://' for asyncpg test")
    print(f"  → Fix your .env: change DATABASE_URL to start with 'postgresql+asyncpg://'")

# ── 1. DNS resolution ─────────────────────────────────────────────────────────
print("\n── Step 1: DNS resolution ─────────────────────────────────────────────")
host = parsed.hostname
try:
    ip = socket.gethostbyname(host)
    print(f"  OK    {host} → {ip}")
except socket.gaierror as e:
    print(f"  FAIL  Cannot resolve {host}: {e}")
    sys.exit(1)

# ── 2. TCP reachability ───────────────────────────────────────────────────────
print("\n── Step 2: TCP connect (port 5432) ────────────────────────────────────")
port = parsed.port or 5432
try:
    sock = socket.create_connection((host, port), timeout=10)
    sock.close()
    print(f"  OK    TCP {host}:{port} reachable")
except (socket.timeout, ConnectionRefusedError, OSError) as e:
    print(f"  FAIL  Cannot TCP-connect to {host}:{port}: {e}")
    print("        Common cause: corporate firewall blocks port 5432.")
    print("        Try Neon's pooler port 6543 or use a VPN/hotspot.")
    sys.exit(1)

# ── 3a. TLS handshake (pure Python ssl — no asyncpg, no asyncio) ──────────────
print("\n── Step 3a: TLS handshake (no asyncpg) ────────────────────────────────")
import ssl as ssl_mod
import threading

def _tls_test(result):
    try:
        ctx = ssl_mod.SSLContext(ssl_mod.PROTOCOL_TLS_CLIENT)
        ctx.check_hostname = False
        ctx.verify_mode = ssl_mod.CERT_NONE
        raw = socket.create_connection((parsed.hostname, parsed.port or 5432), timeout=10)
        tls = ctx.wrap_socket(raw, server_hostname=parsed.hostname)
        result["ok"] = True
        result["ver"] = tls.version()
        tls.close()
    except Exception as e:
        result["err"] = f"{type(e).__name__}: {e}"

_r = {}
_t = threading.Thread(target=_tls_test, args=(_r,), daemon=True)
_t.start()
_t.join(timeout=12)
if _t.is_alive():
    print("  FAIL  TLS timed out — SSL port blocked or Neon not responding to TLS")
    sys.exit(1)
elif "err" in _r:
    print(f"  FAIL  TLS error: {_r['err']}")
    sys.exit(1)
else:
    print(f"  OK    TLS handshake succeeded, protocol={_r['ver']}")

# ── 3b. psycopg (psycopg3) sync connection — no asyncio, no Windows issues ────
print("\n── Step 3b: psycopg connection ────────────────────────────────────────")
try:
    import psycopg
except ImportError:
    print("  FAIL  psycopg not installed — run: uv add 'psycopg[binary]'")
    sys.exit(1)

_host   = parsed.hostname
_port   = parsed.port or 5432
_user   = parsed.username
_passw  = parsed.password
_dbname = parsed.path.lstrip("/")

def _psycopg_thread(result):
    try:
        conn = psycopg.connect(
            host=_host, port=_port, user=_user,
            password=_passw, dbname=_dbname,
            sslmode="require", connect_timeout=30,
        )
        row = conn.execute("SELECT version()").fetchone()
        conn.close()
        result["version"] = row[0]
    except Exception as e:
        result["err"] = f"{type(e).__name__}: {e}"

import time as _time
_r2 = {}
_t2 = threading.Thread(target=_psycopg_thread, args=(_r2,), daemon=True)
_t2.start()
print("  ...  waiting", end="", flush=True)
for _i in range(35):
    if not _t2.is_alive():
        break
    _time.sleep(1)
    print(f" {_i+1}s", end="", flush=True)
print()

if _t2.is_alive():
    print("  FAIL  timed out after 35s — check Neon compute status")
    sys.exit(1)
elif "err" in _r2:
    print(f"  FAIL  {_r2['err']}")
    sys.exit(1)
else:
    print(f"  OK    {_r2['version'][:80]}")

# ── 4. SQLAlchemy async engine (psycopg, main-thread asyncio) ─────────────────
print("\n── Step 4: SQLAlchemy async engine ────────────────────────────────────")
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import text as sa_text

sa_url = (f"postgresql+psycopg://{parsed.username}:{parsed.password}"
          f"@{parsed.hostname}:{parsed.port or 5432}{parsed.path}")

async def _sqlalchemy_async():
    eng = create_async_engine(
        sa_url,
        connect_args={"sslmode": "require", "connect_timeout": 30},
        pool_pre_ping=True,
    )
    async with eng.connect() as conn:
        row = (await conn.execute(sa_text("SELECT version()"))).fetchone()
    await eng.dispose()
    return row[0]

print("  ...  connecting (connect_timeout=30s)", flush=True)
try:
    val = asyncio.run(_sqlalchemy_async())
    print(f"  OK    {val[:80]}")
except Exception as e:
    print(f"  FAIL  {type(e).__name__}: {e}")
    sys.exit(1)

print("\n══ All checks passed — Neon database is reachable and ready ══════════\n")