File size: 1,864 Bytes
23cdeed | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | # -*- coding: utf-8 -*-
"""
Shared lazy PostgreSQL helpers.
Importing this module does not require PostgreSQL or psycopg2. A connection is
attempted only when a caller explicitly asks for one.
"""
from __future__ import annotations
import os
def _get_connection():
"""Return a PostgreSQL connection, creating schema on first use."""
database_url = os.getenv("DATABASE_URL", "").strip()
if not database_url:
raise EnvironmentError("DATABASE_URL is not set")
try:
import psycopg2
except Exception as exc:
raise EnvironmentError("psycopg2 is required for PostgreSQL session memory") from exc
conn = psycopg2.connect(database_url)
_ensure_schema(conn)
return conn
def _ensure_schema(conn) -> None:
with conn.cursor() as cur:
cur.execute(
"""
CREATE TABLE IF NOT EXISTS session_memory (
session_id TEXT PRIMARY KEY,
doc_id TEXT NOT NULL,
created_at TIMESTAMP DEFAULT NOW(),
compressed_json JSONB NOT NULL,
raw_path TEXT
);
"""
)
cur.execute(
"""
CREATE TABLE IF NOT EXISTS response_signals (
id SERIAL PRIMARY KEY,
session_id TEXT,
query_hash TEXT,
signal_type TEXT,
created_at TIMESTAMP DEFAULT NOW()
);
"""
)
cur.execute(
"""
CREATE TABLE IF NOT EXISTS session_graph (
source_session TEXT,
target_session TEXT,
confidence FLOAT,
reason TEXT,
created_at TIMESTAMP DEFAULT NOW(),
PRIMARY KEY (source_session, target_session)
);
"""
)
conn.commit()
|