kussssh/IPO-Analyzer
0
1"""2Postgres (Supabase) persistence layer — drop-in replacement for database.py.3 4Uses psycopg v3 with SUPABASE_DB_URL. Every public function has the exact5same signature as its database.py counterpart so database.py can re-export6them transparently when USE_SUPABASE_RUNTIME is True.7"""8from __future__ import annotations9 10import contextlib11import json12from datetime import datetime, timezone13from pathlib import Path14from typing import Any, Dict, Iterable, List, Optional15 16import psycopg17import psycopg.rows18from psycopg_pool import ConnectionPool19 20from backend.config import SUPABASE_DB_URL, RETRIEVER_CACHE_MAX_ITEMS21 22 23# ── Connection pool ───────────────────────────────────────────────────────────24# Keeps 2–10 live TCP connections to Supabase/PgBouncer, eliminating the25# ~200-500 ms handshake overhead that occurred on every single query.26# prepare_threshold=None disables server-side prepared statements, which are27# incompatible with PgBouncer's transaction-mode pooling (port 6543).28 29_pool: Optional[ConnectionPool] = None30 31 32def _get_pool() -> ConnectionPool:33 global _pool34 if _pool is None:35 _pool = ConnectionPool(36 SUPABASE_DB_URL,37 min_size=2,38 max_size=10,39 kwargs={40 "row_factory": psycopg.rows.dict_row,41 "prepare_threshold": None,42 },43 open=True,44 )45 return _pool46 47 48# ── Connection helpers ────────────────────────────────────────────────────────49 50def utc_now() -> str:51 return datetime.now(timezone.utc).isoformat()52 53 54@contextlib.contextmanager55def _pg():56 """Check out a pooled connection, commit on success, rollback on error."""57 with _get_pool().connection() as conn:58 try:59 yield conn60 conn.commit()61 except Exception:62 conn.rollback()63 raise64 65 66# Expose a get_connection alias so any api.py code that imports it still works.67get_connection = _pg68 69 70# ── JSON helpers ──────────────────────────────────────────────────────────────71 72def _load_json(value: Any) -> Dict[str, Any]:73 if not value:74 return {}75 if isinstance(value, dict):76 return value77 try:78 parsed = json.loads(value)79 return parsed if isinstance(parsed, dict) else {}80 except (json.JSONDecodeError, TypeError):81 return {}82 83 84def _dump_json(value: Any) -> str:85 return json.dumps(value, default=str)86 87 88# ── Date helpers (copied from database.py) ────────────────────────────────────89 90def _parse_sebi_date(value: Optional[str]) -> Optional[datetime]:91 raw = (value or "").strip()92 if not raw:93 return None94 for fmt in ("%b %d, %Y", "%B %d, %Y", "%d %b %Y"):95 try:96 return datetime.strptime(raw, fmt)97 except ValueError:98 continue99 return None100 101 102def _looks_like_date(value: Optional[str]) -> bool:103 return _parse_sebi_date(value) is not None104 105 106# ── Company row → payload (copied from database.py) ──────────────────────────107 108def _company_row_to_payload(109 row: Optional[Dict[str, Any]], use_listing_name: bool = False110) -> Optional[Dict[str, Any]]:111 if not row:112 return None113 114 display_name = (115 row.get("listing_name")116 if use_listing_name and row.get("listing_name")117 else row["company_name"]118 )119 listing_date = (120 row.get("listing_date") if use_listing_name else row.get("drhp_date")121 )122 listing_url = (123 row.get("listing_detail_url")124 if use_listing_name and row.get("listing_detail_url")125 else row["drhp_detail_url"]126 )127 128 drhp = {129 "name": display_name,130 "date": listing_date or "",131 "detail_url": listing_url,132 "is_pdf": str(listing_url).lower().endswith(".pdf"),133 }134 rhp = None135 if row.get("rhp_detail_url"):136 rhp = {137 "name": row.get("rhp_name") or row["company_name"],138 "date": row.get("rhp_date") or "",139 "detail_url": row["rhp_detail_url"],140 "is_pdf": str(row["rhp_detail_url"]).lower().endswith(".pdf"),141 "match_score": row.get("rhp_match_score"),142 }143 prospectus = None144 if row.get("prospectus_detail_url"):145 prospectus = {146 "name": row.get("prospectus_name") or row["company_name"],147 "date": row.get("prospectus_date") or "",148 "detail_url": row["prospectus_detail_url"],149 "is_pdf": str(row["prospectus_detail_url"]).lower().endswith(".pdf"),150 "match_score": row.get("prospectus_match_score"),151 }152 153 stage = "drhp"154 latest_date = drhp["date"]155 if prospectus:156 stage = "prospectus"157 latest_date = prospectus.get("date") or latest_date158 elif rhp:159 stage = "rhp"160 latest_date = rhp.get("date") or latest_date161 162 return {163 "id": row["id"],164 "listing_id": row.get("listing_id"),165 "company_name": display_name,166 "table_title": display_name,167 "table_date": listing_date or "",168 "normalized_name": row["normalized_name"],169 "drhp": drhp,170 "rhp": rhp,171 "prospectus": prospectus,172 "has_drhp": True,173 "has_rhp": bool(rhp),174 "has_prospectus": bool(prospectus),175 "stage": stage,176 "latest_date": listing_date or latest_date,177 }178 179 180# ── Schema init ───────────────────────────────────────────────────────────────181 182def init_db() -> None:183 """Create all tables if they don't exist, then verify connection."""184 try:185 with _pg() as conn:186 conn.execute("""187 CREATE TABLE IF NOT EXISTS companies (188 id BIGSERIAL PRIMARY KEY,189 company_name TEXT NOT NULL,190 normalized_name TEXT NOT NULL UNIQUE,191 source_name TEXT,192 drhp_date TEXT,193 drhp_detail_url TEXT NOT NULL,194 last_seen_at TEXT NOT NULL,195 created_at TEXT NOT NULL,196 updated_at TEXT NOT NULL,197 rhp_name TEXT,198 rhp_date TEXT,199 rhp_detail_url TEXT,200 rhp_match_score REAL,201 rhp_checked_at TEXT,202 prospectus_name TEXT,203 prospectus_date TEXT,204 prospectus_detail_url TEXT,205 prospectus_match_score REAL,206 prospectus_checked_at TEXT207 )208 """)209 conn.execute("""210 CREATE TABLE IF NOT EXISTS drhp_listings (211 id BIGSERIAL PRIMARY KEY,212 company_id BIGINT NOT NULL REFERENCES companies(id) ON DELETE CASCADE,213 listing_name TEXT NOT NULL,214 listing_date TEXT,215 detail_url TEXT NOT NULL UNIQUE,216 created_at TEXT NOT NULL,217 updated_at TEXT NOT NULL218 )219 """)220 conn.execute("""221 CREATE TABLE IF NOT EXISTS documents (222 id BIGSERIAL PRIMARY KEY,223 company_id BIGINT NOT NULL REFERENCES companies(id) ON DELETE CASCADE,224 doc_type TEXT NOT NULL,225 detail_url TEXT,226 pdf_url TEXT,227 local_path TEXT,228 pdf_sha256 TEXT,229 available INTEGER NOT NULL DEFAULT 0,230 downloaded_at TEXT,231 processed_at TEXT,232 page_count INTEGER,233 status TEXT NOT NULL DEFAULT 'pending',234 metadata_json TEXT NOT NULL DEFAULT '{}',235 UNIQUE(company_id, doc_type)236 )237 """)238 conn.execute("""239 CREATE TABLE IF NOT EXISTS document_chunks (240 id BIGSERIAL PRIMARY KEY,241 company_id BIGINT NOT NULL REFERENCES companies(id) ON DELETE CASCADE,242 doc_type TEXT NOT NULL,243 chunk_hash TEXT NOT NULL,244 content TEXT NOT NULL,245 section TEXT,246 chunk_type TEXT,247 page INTEGER,248 chunk_index INTEGER,249 is_delta INTEGER NOT NULL DEFAULT 0,250 change_type TEXT NOT NULL DEFAULT 'unchanged',251 metadata_json TEXT NOT NULL DEFAULT '{}',252 UNIQUE(company_id, doc_type, chunk_hash)253 )254 """)255 conn.execute("""256 CREATE TABLE IF NOT EXISTS diffs (257 id BIGSERIAL PRIMARY KEY,258 company_id BIGINT NOT NULL REFERENCES companies(id) ON DELETE CASCADE,259 from_doc_type TEXT NOT NULL,260 to_doc_type TEXT NOT NULL,261 summary_json TEXT NOT NULL,262 report_markdown TEXT NOT NULL DEFAULT '',263 created_at TEXT NOT NULL,264 updated_at TEXT NOT NULL,265 UNIQUE(company_id, from_doc_type, to_doc_type)266 )267 """)268 conn.execute("""269 CREATE TABLE IF NOT EXISTS analyses (270 id BIGSERIAL PRIMARY KEY,271 company_id BIGINT NOT NULL REFERENCES companies(id) ON DELETE CASCADE,272 doc_type TEXT NOT NULL,273 pdf_sha256 TEXT NOT NULL,274 status TEXT NOT NULL DEFAULT 'complete',275 result_json TEXT NOT NULL,276 created_at TEXT NOT NULL,277 updated_at TEXT NOT NULL,278 UNIQUE(company_id, doc_type, pdf_sha256)279 )280 """)281 conn.execute("""282 CREATE TABLE IF NOT EXISTS embedding_cache (283 id BIGSERIAL PRIMARY KEY,284 model_name TEXT NOT NULL,285 chunk_hash TEXT NOT NULL,286 vector_json TEXT NOT NULL,287 dimensions INTEGER NOT NULL DEFAULT 0,288 created_at TEXT NOT NULL,289 updated_at TEXT NOT NULL,290 last_used_at TEXT NOT NULL,291 UNIQUE(model_name, chunk_hash)292 )293 """)294 conn.execute("""295 CREATE TABLE IF NOT EXISTS company_sectors (296 company_name TEXT PRIMARY KEY,297 sector TEXT NOT NULL298 )299 """)300 conn.execute("""301 CREATE TABLE IF NOT EXISTS users (302 id BIGSERIAL PRIMARY KEY,303 email TEXT UNIQUE NOT NULL,304 password_hash TEXT,305 full_name TEXT,306 phone TEXT,307 auth_provider TEXT DEFAULT 'local',308 is_subscribed INTEGER NOT NULL DEFAULT 0,309 analyses_count INTEGER NOT NULL DEFAULT 0,310 subscribed_at TEXT,311 referral_code TEXT UNIQUE,312 referred_by BIGINT,313 created_at TEXT NOT NULL314 )315 """)316 conn.execute("""317 CREATE TABLE IF NOT EXISTS telemetry (318 id BIGSERIAL PRIMARY KEY,319 user_id BIGINT REFERENCES users(id) ON DELETE SET NULL,320 event_type TEXT NOT NULL,321 metadata TEXT,322 created_at TEXT NOT NULL323 )324 """)325 conn.execute("""326 CREATE TABLE IF NOT EXISTS email_otps (327 id BIGSERIAL PRIMARY KEY,328 email TEXT NOT NULL,329 otp_code TEXT NOT NULL,330 expires_at TEXT NOT NULL,331 verified INTEGER NOT NULL DEFAULT 0,332 created_at TEXT NOT NULL333 )334 """)335 conn.execute("""336 CREATE TABLE IF NOT EXISTS referrals (337 id BIGSERIAL PRIMARY KEY,338 referrer_id BIGINT NOT NULL REFERENCES users(id),339 referred_user_id BIGINT NOT NULL UNIQUE REFERENCES users(id),340 qualified INTEGER NOT NULL DEFAULT 0,341 created_at TEXT NOT NULL342 )343 """)344 conn.execute("""345 CREATE TABLE IF NOT EXISTS notifications (346 id BIGSERIAL PRIMARY KEY,347 user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,348 title TEXT NOT NULL,349 body TEXT NOT NULL,350 type TEXT NOT NULL DEFAULT 'system',351 is_read INTEGER NOT NULL DEFAULT 0,352 created_at TEXT NOT NULL353 )354 """)355 print("[SupabaseDB] Schema ready — all tables exist.")356 except Exception as exc:357 print(f"[SupabaseDB] ⚠ init_db failed: {exc}")358 raise359 360 361def _migrate_users_table() -> None:362 """No-op in Supabase mode — schema migrations happen via Supabase dashboard."""363 pass364 365 366# ── Companies ─────────────────────────────────────────────────────────────────367 368def upsert_companies(entries: Iterable[Dict[str, Any]]) -> int:369 now = utc_now()370 rows = list(entries)371 with _pg() as conn:372 valid_names = [e["normalized_name"] for e in rows if e.get("normalized_name")]373 if valid_names:374 ph = ",".join(["%s"] * len(valid_names))375 to_delete = conn.execute(376 f"SELECT id FROM companies WHERE normalized_name NOT IN ({ph})",377 valid_names,378 ).fetchall()379 if to_delete:380 ids = [r["id"] for r in to_delete]381 id_ph = ",".join(["%s"] * len(ids))382 for tbl in ["document_chunks", "diffs", "analyses", "documents", "drhp_listings"]:383 conn.execute(f"DELETE FROM {tbl} WHERE company_id IN ({id_ph})", ids)384 conn.execute(f"DELETE FROM companies WHERE id IN ({id_ph})", ids)385 386 conn.execute("DELETE FROM drhp_listings")387 388 for entry in rows:389 existing_row = conn.execute(390 """391 SELECT rhp_name, rhp_date, rhp_detail_url, rhp_match_score, rhp_checked_at,392 prospectus_name, prospectus_date, prospectus_detail_url,393 prospectus_match_score, prospectus_checked_at394 FROM companies WHERE normalized_name = %s395 """,396 (entry["normalized_name"],),397 ).fetchone()398 399 rhp_name = entry.get("rhp_name")400 rhp_date = entry.get("rhp_date")401 rhp_detail_url = entry.get("rhp_detail_url")402 rhp_match_score = entry.get("rhp_match_score")403 rhp_checked_at = now if rhp_detail_url else None404 if entry.get("preserve_existing_rhp") and existing_row:405 rhp_name = existing_row.get("rhp_name")406 rhp_date = existing_row.get("rhp_date")407 rhp_detail_url = existing_row.get("rhp_detail_url")408 rhp_match_score = existing_row.get("rhp_match_score")409 rhp_checked_at = existing_row.get("rhp_checked_at")410 411 prospectus_name = entry.get("prospectus_name")412 prospectus_date = entry.get("prospectus_date")413 prospectus_detail_url = entry.get("prospectus_detail_url")414 prospectus_match_score = entry.get("prospectus_match_score")415 prospectus_checked_at = now if prospectus_detail_url else None416 if entry.get("preserve_existing_prospectus") and existing_row:417 prospectus_name = existing_row.get("prospectus_name")418 prospectus_date = existing_row.get("prospectus_date")419 prospectus_detail_url = existing_row.get("prospectus_detail_url")420 prospectus_match_score = existing_row.get("prospectus_match_score")421 prospectus_checked_at = existing_row.get("prospectus_checked_at")422 423 conn.execute(424 """425 INSERT INTO companies (426 company_name, normalized_name, source_name, drhp_date, drhp_detail_url,427 last_seen_at, created_at, updated_at,428 rhp_name, rhp_date, rhp_detail_url, rhp_match_score, rhp_checked_at,429 prospectus_name, prospectus_date, prospectus_detail_url,430 prospectus_match_score, prospectus_checked_at431 ) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)432 ON CONFLICT(normalized_name) DO UPDATE SET433 company_name=EXCLUDED.company_name,434 source_name=EXCLUDED.source_name,435 drhp_date=EXCLUDED.drhp_date,436 drhp_detail_url=EXCLUDED.drhp_detail_url,437 last_seen_at=EXCLUDED.last_seen_at,438 rhp_name=EXCLUDED.rhp_name,439 rhp_date=EXCLUDED.rhp_date,440 rhp_detail_url=EXCLUDED.rhp_detail_url,441 rhp_match_score=EXCLUDED.rhp_match_score,442 rhp_checked_at=EXCLUDED.rhp_checked_at,443 prospectus_name=EXCLUDED.prospectus_name,444 prospectus_date=EXCLUDED.prospectus_date,445 prospectus_detail_url=EXCLUDED.prospectus_detail_url,446 prospectus_match_score=EXCLUDED.prospectus_match_score,447 prospectus_checked_at=EXCLUDED.prospectus_checked_at,448 updated_at=EXCLUDED.updated_at449 """,450 (451 entry["company_name"],452 entry["normalized_name"],453 entry.get("source_name", entry["company_name"]),454 entry.get("date", ""),455 entry["detail_url"],456 now, now, now,457 rhp_name, rhp_date, rhp_detail_url, rhp_match_score, rhp_checked_at,458 prospectus_name, prospectus_date, prospectus_detail_url,459 prospectus_match_score, prospectus_checked_at,460 ),461 )462 463 company_row = conn.execute(464 "SELECT id FROM companies WHERE normalized_name = %s",465 (entry["normalized_name"],),466 ).fetchone()467 if not company_row:468 continue469 470 conn.execute(471 """472 INSERT INTO drhp_listings (473 company_id, listing_name, listing_date, detail_url, created_at, updated_at474 ) VALUES (%s,%s,%s,%s,%s,%s)475 ON CONFLICT(detail_url) DO UPDATE SET476 company_id=EXCLUDED.company_id,477 listing_name=EXCLUDED.listing_name,478 listing_date=EXCLUDED.listing_date,479 updated_at=EXCLUDED.updated_at480 """,481 (482 company_row["id"],483 entry.get("source_name", entry["company_name"]),484 entry.get("date", ""),485 entry["detail_url"],486 now, now,487 ),488 )489 490 conn.execute(491 """492 INSERT INTO documents (company_id, doc_type, detail_url, available, status, metadata_json)493 SELECT id, 'drhp', %s, 1, 'cataloged', '{}'494 FROM companies WHERE normalized_name = %s495 ON CONFLICT(company_id, doc_type) DO UPDATE SET496 detail_url=EXCLUDED.detail_url, available=1, status='cataloged'497 """,498 (entry["detail_url"], entry["normalized_name"]),499 )500 501 for doc_type in ("rhp", "prospectus"):502 d_url = rhp_detail_url if doc_type == "rhp" else prospectus_detail_url503 conn.execute(504 """505 INSERT INTO documents (company_id, doc_type, detail_url, available, status, metadata_json)506 VALUES (%s,%s,%s,%s,%s,'{}')507 ON CONFLICT(company_id, doc_type) DO UPDATE SET508 detail_url=EXCLUDED.detail_url,509 available=EXCLUDED.available,510 status=EXCLUDED.status511 """,512 (513 company_row["id"],514 doc_type,515 d_url,516 1 if d_url else 0,517 "cataloged" if d_url else "missing",518 ),519 )520 return len(rows)521 522 523def purge_invalid_catalog_rows() -> int:524 with _pg() as conn:525 listing_rows = conn.execute(526 "SELECT id, listing_name, listing_date FROM drhp_listings"527 ).fetchall()528 bad_listing_ids = [529 r["id"]530 for r in listing_rows531 if _looks_like_date(r.get("listing_name"))532 or (r.get("listing_date") and not _looks_like_date(r.get("listing_date")))533 ]534 if bad_listing_ids:535 ph = ",".join(["%s"] * len(bad_listing_ids))536 conn.execute(f"DELETE FROM drhp_listings WHERE id IN ({ph})", bad_listing_ids)537 538 company_rows = conn.execute(539 "SELECT id, company_name, drhp_date, normalized_name FROM companies"540 ).fetchall()541 bad_company_ids = [542 r["id"]543 for r in company_rows544 if _looks_like_date(r.get("company_name"))545 or (r.get("drhp_date") and not _looks_like_date(r.get("drhp_date")))546 ]547 if bad_company_ids:548 ph = ",".join(["%s"] * len(bad_company_ids))549 conn.execute(f"DELETE FROM companies WHERE id IN ({ph})", bad_company_ids)550 551 return len(bad_listing_ids) + len(bad_company_ids)552 553 554def list_companies() -> List[Dict[str, Any]]:555 with _pg() as conn:556 listing_rows = conn.execute(557 """558 SELECT l.id AS listing_id, l.listing_name, l.listing_date,559 l.detail_url AS listing_detail_url, c.*560 FROM drhp_listings l561 JOIN companies c ON c.id = l.company_id562 ORDER BY l.id DESC563 """564 ).fetchall()565 if listing_rows:566 listing_rows = sorted(567 listing_rows,568 key=lambda r: (569 _parse_sebi_date(r.get("listing_date")) or datetime.min,570 (r.get("listing_name") or "").lower(),571 r.get("listing_id", 0),572 ),573 reverse=True,574 )575 return [_company_row_to_payload(r, use_listing_name=True) for r in listing_rows]576 return []577 578 579def get_company(company_id: int) -> Optional[Dict[str, Any]]:580 with _pg() as conn:581 row = conn.execute(582 """583 SELECT l.id AS listing_id, l.listing_name, l.listing_date,584 l.detail_url AS listing_detail_url, c.*585 FROM companies c586 LEFT JOIN drhp_listings l ON l.company_id = c.id587 WHERE c.id = %s588 ORDER BY589 CASE WHEN l.listing_date IS NULL THEN 1 ELSE 0 END,590 l.listing_date DESC, l.id DESC591 LIMIT 1592 """,593 (company_id,),594 ).fetchone()595 return _company_row_to_payload(row, use_listing_name=bool(row and row.get("listing_name"))) if row else None596 597 598def get_company_row(company_id: int) -> Optional[Dict[str, Any]]:599 with _pg() as conn:600 return conn.execute(601 "SELECT * FROM companies WHERE id = %s", (company_id,)602 ).fetchone()603 604 605def update_company_document_match(606 company_id: int,607 doc_type: str,608 match: Optional[Dict[str, Any]],609 checked_at: Optional[str] = None,610) -> None:611 checked_at = checked_at or utc_now()612 if doc_type not in {"rhp", "prospectus"}:613 raise ValueError(f"Unsupported doc_type: {doc_type}")614 615 name_col = f"{doc_type}_name"616 date_col = f"{doc_type}_date"617 url_col = f"{doc_type}_detail_url"618 score_col = f"{doc_type}_match_score"619 checked_col = f"{doc_type}_checked_at"620 621 with _pg() as conn:622 conn.execute(623 f"""624 UPDATE companies625 SET {name_col}=%s, {date_col}=%s, {url_col}=%s,626 {score_col}=%s, {checked_col}=%s, updated_at=%s627 WHERE id=%s628 """,629 (630 match.get("name") if match else None,631 match.get("date") if match else None,632 match.get("detail_url") if match else None,633 match.get("match_score") if match else None,634 checked_at,635 checked_at,636 company_id,637 ),638 )639 conn.execute(640 """641 INSERT INTO documents (company_id, doc_type, detail_url, available, status, metadata_json)642 VALUES (%s,%s,%s,%s,%s,'{}')643 ON CONFLICT(company_id, doc_type) DO UPDATE SET644 detail_url=EXCLUDED.detail_url,645 available=EXCLUDED.available,646 status=EXCLUDED.status647 """,648 (649 company_id,650 doc_type,651 match.get("detail_url") if match else None,652 1 if match else 0,653 "cataloged" if match else "missing",654 ),655 )656 657 658# ── Documents ─────────────────────────────────────────────────────────────────659 660def update_document_record(661 company_id: int,662 doc_type: str,663 *,664 detail_url: Optional[str] = None,665 pdf_url: Optional[str] = None,666 local_path: Optional[str] = None,667 pdf_sha256: Optional[str] = None,668 available: Optional[bool] = None,669 downloaded_at: Optional[str] = None,670 processed_at: Optional[str] = None,671 page_count: Optional[int] = None,672 status: Optional[str] = None,673 metadata: Optional[Dict[str, Any]] = None,674 merge_metadata: Optional[Dict[str, Any]] = None,675) -> None:676 existing = get_document_record(company_id, doc_type) or {}677 existing_metadata = _load_json(existing.get("metadata_json"))678 if metadata is not None:679 metadata_payload = metadata680 elif merge_metadata:681 metadata_payload = {**existing_metadata, **merge_metadata}682 else:683 metadata_payload = existing_metadata684 685 payload = {686 "detail_url": detail_url if detail_url is not None else existing.get("detail_url"),687 "pdf_url": pdf_url if pdf_url is not None else existing.get("pdf_url"),688 "local_path": local_path if local_path is not None else existing.get("local_path"),689 "pdf_sha256": pdf_sha256 if pdf_sha256 is not None else existing.get("pdf_sha256"),690 "available": int(existing.get("available", 0) if available is None else available),691 "downloaded_at": downloaded_at if downloaded_at is not None else existing.get("downloaded_at"),692 "processed_at": processed_at if processed_at is not None else existing.get("processed_at"),693 "page_count": page_count if page_count is not None else existing.get("page_count"),694 "status": status if status is not None else existing.get("status", "pending"),695 "metadata_json": _dump_json(metadata_payload),696 }697 698 with _pg() as conn:699 conn.execute(700 """701 INSERT INTO documents (702 company_id, doc_type, detail_url, pdf_url, local_path, pdf_sha256,703 available, downloaded_at, processed_at, page_count, status, metadata_json704 ) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)705 ON CONFLICT(company_id, doc_type) DO UPDATE SET706 detail_url=EXCLUDED.detail_url,707 pdf_url=EXCLUDED.pdf_url,708 local_path=EXCLUDED.local_path,709 pdf_sha256=EXCLUDED.pdf_sha256,710 available=EXCLUDED.available,711 downloaded_at=EXCLUDED.downloaded_at,712 processed_at=EXCLUDED.processed_at,713 page_count=EXCLUDED.page_count,714 status=EXCLUDED.status,715 metadata_json=EXCLUDED.metadata_json716 """,717 (718 company_id, doc_type,719 payload["detail_url"], payload["pdf_url"], payload["local_path"],720 payload["pdf_sha256"], payload["available"], payload["downloaded_at"],721 payload["processed_at"], payload["page_count"],722 payload["status"], payload["metadata_json"],723 ),724 )725 726 727def get_document_record(company_id: int, doc_type: str) -> Optional[Dict[str, Any]]:728 with _pg() as conn:729 return conn.execute(730 "SELECT * FROM documents WHERE company_id=%s AND doc_type=%s",731 (company_id, doc_type),732 ).fetchone()733 734 735# ── Chunks ────────────────────────────────────────────────────────────────────736 737def clear_document_chunks(company_id: int, doc_type: str) -> None:738 with _pg() as conn:739 conn.execute(740 "DELETE FROM document_chunks WHERE company_id=%s AND doc_type=%s",741 (company_id, doc_type),742 )743 744 745def save_document_chunks(company_id: int, doc_type: str, chunks: List[Dict[str, Any]]) -> None:746 deduped_chunks: List[Dict[str, Any]] = []747 seen_hashes: Dict[str, Dict[str, Any]] = {}748 749 for chunk in chunks:750 chunk_hash = chunk["chunk_hash"]751 existing = seen_hashes.get(chunk_hash)752 if not existing:753 metadata = dict(chunk.get("metadata", {}))754 metadata["duplicate_pages"] = [chunk.get("page")]755 metadata["duplicate_chunk_indexes"] = [chunk.get("chunk_index")]756 deduped = {**chunk, "metadata": metadata}757 seen_hashes[chunk_hash] = deduped758 deduped_chunks.append(deduped)759 continue760 metadata = existing.setdefault("metadata", {})761 pages = metadata.setdefault("duplicate_pages", [])762 indexes = metadata.setdefault("duplicate_chunk_indexes", [])763 if chunk.get("page") not in pages:764 pages.append(chunk.get("page"))765 if chunk.get("chunk_index") not in indexes:766 indexes.append(chunk.get("chunk_index"))767 768 _CHUNK_INSERT_BATCH = 500 # max rows per INSERT; keeps param count < 65535769 770 with _pg() as conn:771 conn.execute(772 "DELETE FROM document_chunks WHERE company_id=%s AND doc_type=%s",773 (company_id, doc_type),774 )775 # Batch multi-row INSERT — reduces 1800+ round-trips to ~4776 for i in range(0, len(deduped_chunks), _CHUNK_INSERT_BATCH):777 batch = deduped_chunks[i : i + _CHUNK_INSERT_BATCH]778 placeholders = ",".join(["(%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)"] * len(batch))779 params: List[Any] = []780 for chunk in batch:781 params.extend([782 company_id, doc_type,783 chunk["chunk_hash"], chunk["content"],784 chunk.get("section"), chunk.get("chunk_type"),785 chunk.get("page"), chunk.get("chunk_index"),786 int(chunk.get("is_delta", False)),787 chunk.get("change_type", "unchanged"),788 _dump_json(chunk.get("metadata", {})),789 ])790 conn.execute(791 f"""792 INSERT INTO document_chunks (793 company_id, doc_type, chunk_hash, content, section, chunk_type,794 page, chunk_index, is_delta, change_type, metadata_json795 ) VALUES {placeholders}796 ON CONFLICT(company_id, doc_type, chunk_hash) DO UPDATE SET797 content=EXCLUDED.content, section=EXCLUDED.section,798 chunk_type=EXCLUDED.chunk_type, page=EXCLUDED.page,799 chunk_index=EXCLUDED.chunk_index, is_delta=EXCLUDED.is_delta,800 change_type=EXCLUDED.change_type, metadata_json=EXCLUDED.metadata_json801 """,802 params,803 )804 805 806def get_document_chunks(company_id: int, doc_type: str) -> List[Dict[str, Any]]:807 with _pg() as conn:808 rows = conn.execute(809 """810 SELECT * FROM document_chunks811 WHERE company_id=%s AND doc_type=%s812 ORDER BY page ASC, chunk_index ASC813 """,814 (company_id, doc_type),815 ).fetchall()816 return [817 {818 "chunk_hash": row["chunk_hash"],819 "content": row["content"],820 "section": row["section"],821 "chunk_type": row["chunk_type"],822 "page": row["page"],823 "chunk_index": row["chunk_index"],824 "is_delta": bool(row["is_delta"]),825 "change_type": row["change_type"],826 "metadata": _load_json(row.get("metadata_json")),827 }828 for row in rows829 ]830 831 832# ── Diffs ─────────────────────────────────────────────────────────────────────833 834def save_diff(company_id: int, from_doc_type: str, to_doc_type: str, diff: Dict[str, Any]) -> None:835 now = utc_now()836 with _pg() as conn:837 conn.execute(838 """839 INSERT INTO diffs (840 company_id, from_doc_type, to_doc_type,841 summary_json, report_markdown, created_at, updated_at842 ) VALUES (%s,%s,%s,%s,%s,%s,%s)843 ON CONFLICT(company_id, from_doc_type, to_doc_type) DO UPDATE SET844 summary_json=EXCLUDED.summary_json,845 report_markdown=EXCLUDED.report_markdown,846 updated_at=EXCLUDED.updated_at847 """,848 (849 company_id, from_doc_type, to_doc_type,850 _dump_json(diff), diff.get("report_markdown", ""),851 now, now,852 ),853 )854 855 856def get_diff(company_id: int, from_doc_type: str, to_doc_type: str) -> Optional[Dict[str, Any]]:857 with _pg() as conn:858 row = conn.execute(859 """860 SELECT summary_json FROM diffs861 WHERE company_id=%s AND from_doc_type=%s AND to_doc_type=%s862 """,863 (company_id, from_doc_type, to_doc_type),864 ).fetchone()865 return _load_json(row["summary_json"]) if row else None866 867 868# ── Analyses ──────────────────────────────────────────────────────────────────869 870def save_analysis(company_id: int, doc_type: str, pdf_sha256: str, result: Dict[str, Any]) -> None:871 now = utc_now()872 with _pg() as conn:873 conn.execute(874 """875 INSERT INTO analyses (876 company_id, doc_type, pdf_sha256, status, result_json, created_at, updated_at877 ) VALUES (%s,%s,%s,'complete',%s,%s,%s)878 ON CONFLICT(company_id, doc_type, pdf_sha256) DO UPDATE SET879 status='complete',880 result_json=EXCLUDED.result_json,881 updated_at=EXCLUDED.updated_at882 """,883 (company_id, doc_type, pdf_sha256, _dump_json(result), now, now),884 )885 886 887def get_analysis(company_id: int, doc_type: str, pdf_sha256: str) -> Optional[Dict[str, Any]]:888 with _pg() as conn:889 row = conn.execute(890 """891 SELECT result_json FROM analyses892 WHERE company_id=%s AND doc_type=%s AND pdf_sha256=%s893 """,894 (company_id, doc_type, pdf_sha256),895 ).fetchone()896 return _load_json(row["result_json"]) if row else None897 898 899def delete_analysis(company_id: int, doc_type: Optional[str] = None) -> int:900 with _pg() as conn:901 if doc_type:902 cur = conn.execute(903 "DELETE FROM analyses WHERE company_id=%s AND doc_type=%s",904 (company_id, doc_type),905 )906 else:907 cur = conn.execute(908 "DELETE FROM analyses WHERE company_id=%s", (company_id,)909 )910 return cur.rowcount911 912 913def reset_analysis_state() -> Dict[str, int]:914 with _pg() as conn:915 a = conn.execute("DELETE FROM analyses").rowcount916 d = conn.execute("DELETE FROM diffs").rowcount917 c = conn.execute("DELETE FROM document_chunks").rowcount918 r = conn.execute(919 """920 UPDATE documents SET921 pdf_url=NULL, local_path=NULL, pdf_sha256=NULL,922 downloaded_at=NULL, processed_at=NULL, page_count=NULL,923 status=CASE924 WHEN detail_url IS NOT NULL AND available=1 THEN 'cataloged'925 WHEN detail_url IS NULL OR available=0 THEN 'missing'926 ELSE 'pending'927 END,928 metadata_json='{}'929 """930 ).rowcount931 return {932 "analyses_deleted": a,933 "diffs_deleted": d,934 "chunks_deleted": c,935 "documents_reset": r,936 }937 938 939def list_cached_analyses() -> List[Dict[str, Any]]:940 with _pg() as conn:941 rows = conn.execute(942 "SELECT result_json FROM analyses ORDER BY updated_at DESC"943 ).fetchall()944 return [_load_json(r["result_json"]) for r in rows]945 946 947# ── Embedding cache ───────────────────────────────────────────────────────────948 949def get_cached_embeddings(model_name: str, chunk_hashes: Iterable[str]) -> Dict[str, List[float]]:950 hashes = [h for h in chunk_hashes if h]951 if not hashes:952 return {}953 ph = ",".join(["%s"] * len(hashes))954 with _pg() as conn:955 rows = conn.execute(956 f"""957 SELECT chunk_hash, vector_json FROM embedding_cache958 WHERE model_name=%s AND chunk_hash IN ({ph})959 """,960 (model_name, *hashes),961 ).fetchall()962 if rows:963 now = utc_now()964 hashes_found = [r["chunk_hash"] for r in rows]965 ph2 = ",".join(["%s"] * len(hashes_found))966 conn.execute(967 f"UPDATE embedding_cache SET last_used_at=%s, updated_at=%s WHERE model_name=%s AND chunk_hash IN ({ph2})",968 (now, now, model_name, *hashes_found),969 )970 cached: Dict[str, List[float]] = {}971 for row in rows:972 vector = _load_json(row.get("vector_json"))973 if isinstance(vector, list):974 cached[row["chunk_hash"]] = vector975 return cached976 977 978def save_cached_embeddings(model_name: str, vectors_by_hash: Dict[str, List[float]]) -> None:979 if not vectors_by_hash:980 return981 now = utc_now()982 valid = [983 (model_name, h, _dump_json(v), len(v), now, now, now)984 for h, v in vectors_by_hash.items()985 if h and isinstance(v, list)986 ]987 if not valid:988 return989 # Batch at 100 rows — vectors are large JSON strings, keep payload manageable990 _EMBED_BATCH = 100991 with _pg() as conn:992 for i in range(0, len(valid), _EMBED_BATCH):993 batch = valid[i : i + _EMBED_BATCH]994 placeholders = ",".join(["(%s,%s,%s,%s,%s,%s,%s)"] * len(batch))995 params: List[Any] = [v for row in batch for v in row]996 conn.execute(997 f"""998 INSERT INTO embedding_cache (999 model_name, chunk_hash, vector_json, dimensions,1000 created_at, updated_at, last_used_at1001 ) VALUES {placeholders}1002 ON CONFLICT(model_name, chunk_hash) DO UPDATE SET1003 vector_json=EXCLUDED.vector_json,1004 dimensions=EXCLUDED.dimensions,1005 updated_at=EXCLUDED.updated_at,1006 last_used_at=EXCLUDED.last_used_at1007 """,1008 params,1009 )1010 1011 1012# ── Sectors ───────────────────────────────────────────────────────────────────1013 1014def get_company_sector(company_name: str) -> Optional[str]:1015 with _pg() as conn:1016 row = conn.execute(1017 "SELECT sector FROM company_sectors WHERE company_name=%s", (company_name,)1018 ).fetchone()1019 return row["sector"] if row else None1020 1021 1022def save_company_sector(company_name: str, sector: str) -> None:1023 with _pg() as conn:1024 conn.execute(1025 """1026 INSERT INTO company_sectors (company_name, sector) VALUES (%s,%s)1027 ON CONFLICT(company_name) DO UPDATE SET sector=EXCLUDED.sector1028 """,1029 (company_name, sector),1030 )1031 1032 1033# ── Users ─────────────────────────────────────────────────────────────────────1034 1035def create_user(1036 email: str,1037 password_hash: str,1038 full_name: str,1039 auth_provider: str = "local",1040 phone: Optional[str] = None,1041 referral_code: Optional[str] = None,1042) -> Optional[int]:1043 from backend.referral_utils import generate_referral_code as _gen_code1044 my_code = _gen_code()1045 try:1046 with _pg() as conn:1047 cur = conn.execute(1048 """1049 INSERT INTO users (1050 email, password_hash, full_name, phone, auth_provider,1051 referral_code, referred_by, created_at1052 ) VALUES (%s,%s,%s,%s,%s,%s,1053 (SELECT id FROM users WHERE referral_code=%s),1054 %s)1055 RETURNING id1056 """,1057 (email, password_hash, full_name, phone, auth_provider,1058 my_code, referral_code, utc_now()),1059 )1060 row = cur.fetchone()1061 return row["id"] if row else None1062 except Exception:1063 return None1064 1065 1066def get_user_by_email(email: str) -> Optional[Dict[str, Any]]:1067 with _pg() as conn:1068 row = conn.execute("SELECT * FROM users WHERE email=%s", (email,)).fetchone()1069 if row and row.get("email") == "kushgupta2019@gmail.com":1070 row = dict(row)1071 row["is_subscribed"] = 11072 return row1073 1074 1075def get_user_by_id(user_id: int) -> Optional[Dict[str, Any]]:1076 with _pg() as conn:1077 row = conn.execute("SELECT * FROM users WHERE id=%s", (user_id,)).fetchone()1078 if row and row.get("email") == "kushgupta2019@gmail.com":1079 row = dict(row)1080 row["is_subscribed"] = 11081 return row1082 1083 1084def log_telemetry(user_id: Optional[int], event_type: str, metadata: Dict[str, Any]) -> None:1085 try:1086 with _pg() as conn:1087 conn.execute(1088 "INSERT INTO telemetry (user_id, event_type, metadata, created_at) VALUES (%s,%s,%s,%s)",1089 (user_id, event_type, _dump_json(metadata), utc_now()),1090 )1091 except Exception as exc:1092 # Surface DB errors in production logs (HF Spaces stdout)1093 print(f"[supabase_db] log_telemetry FAILED user={user_id} event={event_type}: {exc}")1094 raise # re-raise so api.py can catch and log too1095 1096 1097def increment_user_analysis_count(user_id: int) -> int:1098 with _pg() as conn:1099 conn.execute(1100 "UPDATE users SET analyses_count=analyses_count+1 WHERE id=%s", (user_id,)1101 )1102 row = conn.execute(1103 "SELECT analyses_count FROM users WHERE id=%s", (user_id,)1104 ).fetchone()1105 return row["analyses_count"] if row else 01106 1107 1108def upgrade_user_subscription(user_id: int) -> None:1109 with _pg() as conn:1110 conn.execute(1111 "UPDATE users SET is_subscribed=1, subscribed_at=%s WHERE id=%s",1112 (utc_now(), user_id),1113 )1114 1115 1116def update_user_phone(user_id: int, phone: str) -> None:1117 with _pg() as conn:1118 conn.execute("UPDATE users SET phone=%s WHERE id=%s", (phone, user_id))1119 1120 1121# ── OTP ───────────────────────────────────────────────────────────────────────1122 1123def store_otp(email: str, otp_code: str, expires_at: str) -> None:1124 with _pg() as conn:1125 conn.execute("DELETE FROM email_otps WHERE email=%s", (email,))1126 conn.execute(1127 "INSERT INTO email_otps (email, otp_code, expires_at, verified, created_at) VALUES (%s,%s,%s,0,%s)",1128 (email, otp_code, expires_at, utc_now()),1129 )1130 1131 1132def verify_otp(email: str, otp_code: str) -> bool:1133 with _pg() as conn:1134 row = conn.execute(1135 "SELECT id, expires_at, verified FROM email_otps WHERE email=%s AND otp_code=%s",1136 (email, otp_code),1137 ).fetchone()1138 if not row:1139 return False1140 if row["verified"]:1141 return False1142 expires = datetime.fromisoformat(str(row["expires_at"]))1143 if expires.tzinfo is None:1144 expires = expires.replace(tzinfo=timezone.utc)1145 if datetime.now(timezone.utc) > expires:1146 return False1147 conn.execute("UPDATE email_otps SET verified=1 WHERE id=%s", (row["id"],))1148 return True1149 1150 1151def cleanup_expired_otps() -> None:1152 with _pg() as conn:1153 conn.execute("DELETE FROM email_otps WHERE expires_at < %s", (utc_now(),))1154 1155 1156# ── Referrals ─────────────────────────────────────────────────────────────────1157 1158def record_referral(referrer_id: int, referred_user_id: int) -> None:1159 try:1160 with _pg() as conn:1161 conn.execute(1162 """1163 INSERT INTO referrals (referrer_id, referred_user_id, qualified, created_at)1164 VALUES (%s,%s,0,%s)1165 ON CONFLICT(referred_user_id) DO NOTHING1166 """,1167 (referrer_id, referred_user_id, utc_now()),1168 )1169 except Exception:1170 pass1171 1172 1173def qualify_referral(referred_user_id: int) -> Optional[int]:1174 with _pg() as conn:1175 row = conn.execute(1176 "SELECT id, referrer_id, qualified FROM referrals WHERE referred_user_id=%s",1177 (referred_user_id,),1178 ).fetchone()1179 if not row or row["qualified"]:1180 return None1181 conn.execute("UPDATE referrals SET qualified=1 WHERE id=%s", (row["id"],))1182 return row["referrer_id"]1183 1184 1185def check_and_grant_reward(referrer_id: int) -> bool:1186 with _pg() as conn:1187 count_row = conn.execute(1188 "SELECT COUNT(*) AS cnt FROM referrals WHERE referrer_id=%s AND qualified=1",1189 (referrer_id,),1190 ).fetchone()1191 if not count_row or count_row["cnt"] < 5:1192 return False1193 user = conn.execute(1194 "SELECT is_subscribed FROM users WHERE id=%s", (referrer_id,)1195 ).fetchone()1196 if not user or user["is_subscribed"]:1197 return False1198 conn.execute(1199 "UPDATE users SET is_subscribed=1, subscribed_at=%s WHERE id=%s",1200 (utc_now(), referrer_id),