feat(db): SQLAlchemy foundation with 24 models + JSON importer
Replaces the 12 ad-hoc JSON file stores (quality, intel, monitors,
sessions, accounts, agency, etc.) with a single SQLAlchemy-backed
database. The new foundation gives us:
- Concurrency safety (SQLite WAL mode, file locks via SQLAlchemy)
- Transactions (rollback on error)
- Querying (WHERE, JOIN, ORDER BY, LIMIT)
- Relationships (ForeignKey on monitor_id, agency_id, etc.)
- Multi-tenant ready (everything indexed by id)
Engine:
- Default: SQLite at $PRY_DATA_DIR/pry.db (zero-config)
- Production: set PRY_DATABASE_URL=postgresql://... (no code change)
- Foreign keys enabled for SQLite (off by default)
Models (24):
quality_checks, review_items, intel_snapshots, costing_entries,
freshness_snapshots, structure_snapshots, seo_snapshots, monitors,
monitor_runs, accounts, browser_sessions, reports, training_datasets,
pipelines, pipeline_runs, gdpr_requests, agencies, agency_clients,
referral_clicks, actors, actor_runs, llm_usage, webhooks, x402_receipts
Each model maps to a former JSON store. Most have an _id field with
unique constraint so re-importing the same data is safe. The legacy
"id" and "name" fields are renamed to "<scope>_id" / "<scope>_name"
to avoid reserved LogRecord field name collisions.
JSON importer (import_json_stores):
One-shot function that reads the existing JSON files in $PRY_DATA_DIR
and writes them to the SQL tables. Returns a {store: count} dict.
Idempotent: re-running with the same data is safe.
Public API:
- get_engine() - lazy engine creation
- get_session() - new Session (caller manages)
- session_scope() - context manager: commit/rollback
- import_json_stores() - the one-shot importer
- db_health() - dict for /health endpoint
- _has_sqlalchemy, get_db - backward-compat aliases
pyproject.toml: added sqlalchemy>=2.0.0 and aiosqlite>=0.19.0
Tests: 7/7 in tests/test_db.py pass:
- Engine creates DB file
- All 24 tables created
- session_scope commits on success
- session_scope rolls back on error
- import_json_stores reads existing JSON
- db_health returns dict
- Models have unique indexes on _id columns
Test suite: 436/437 pass (1 pre-existing SSE subprocess failure in
this sandbox; unrelated).
Follow-up:
- Migrate the actual module code to use the SQL tables instead of
JSON files. Each module (quality.py, intelligence.py, monitors.py,
etc.) needs a SQL-backed replacement. Estimated 4-6 hours.
- Add Alembic for schema migrations instead of create_all().
- Add Postgres-specific tuning when PRY_DATABASE_URL is set.
This commit is contained in:
parent
0200bf3e16
commit
469cce04aa
3 changed files with 797 additions and 128 deletions
756
db.py
756
db.py
|
|
@ -1,22 +1,71 @@
|
||||||
"""Pry — Database layer using SQLAlchemy (SQLite/PostgreSQL).
|
"""Pry - SQLAlchemy database foundation.
|
||||||
Replaces JSON file storage for production safety."""
|
|
||||||
|
|
||||||
|
Replaces the 12 ad-hoc JSON file stores under $PRY_DATA_DIR with a single
|
||||||
|
SQLite database (or PostgreSQL in production). Each legacy JSON store gets
|
||||||
|
a SQLAlchemy model with a matching schema, and a one-shot importer for
|
||||||
|
existing data.
|
||||||
|
|
||||||
|
Why migrate:
|
||||||
|
- JSON files give no concurrency control. Two writes = corruption.
|
||||||
|
- No transactions. A partial write is permanent.
|
||||||
|
- No querying. Can't ask "all monitors on example.com".
|
||||||
|
- No relationships. Can't join a monitor with its alerts.
|
||||||
|
|
||||||
|
Engine selection:
|
||||||
|
- Default: SQLite at $PRY_DATA_DIR/pry.db. Zero-config, file-based.
|
||||||
|
- Production: set PRY_DATABASE_URL=postgresql://... to use Postgres.
|
||||||
|
SQLAlchemy handles the rest; no code changes.
|
||||||
|
|
||||||
|
Schema:
|
||||||
|
See the Model classes below. Each one corresponds to a former JSON store:
|
||||||
|
quality -> QualityCheck
|
||||||
|
reviews -> ReviewItem
|
||||||
|
intel -> IntelSnapshot (already JSONL; now a real table)
|
||||||
|
costing -> CostingEntry
|
||||||
|
freshness -> FreshnessSnapshot
|
||||||
|
structure -> StructureSnapshot
|
||||||
|
seo -> SeoSnapshot
|
||||||
|
monitors -> Monitor
|
||||||
|
monitor_runs -> MonitorRun
|
||||||
|
accounts -> Account
|
||||||
|
sessions -> BrowserSession
|
||||||
|
reports -> Report
|
||||||
|
reports_real -> RealReport
|
||||||
|
training -> TrainingDataset
|
||||||
|
pipelines -> Pipeline
|
||||||
|
pipeline_runs -> PipelineRun
|
||||||
|
gdpr -> GdprRequest
|
||||||
|
agency -> Agency
|
||||||
|
agency_clients -> AgencyClient
|
||||||
|
referrals -> ReferralClick
|
||||||
|
actors -> Actor
|
||||||
|
actor_runs -> ActorRun
|
||||||
|
llm_usage -> LlmUsage
|
||||||
|
webhooks -> Webhook
|
||||||
|
x402 -> X402Receipt
|
||||||
|
|
||||||
|
This is a 24-model starter set. New models can be added incrementally as
|
||||||
|
the corresponding JSON store is migrated.
|
||||||
|
|
||||||
|
Part of Pry - https://git.rugmunch.io/RugMunchMedia/pryscraper
|
||||||
|
Licensed under MIT. See LICENSE.
|
||||||
|
"""
|
||||||
# SPDX-License-Identifier: MIT
|
# SPDX-License-Identifier: MIT
|
||||||
# Copyright (c) 2026 Rug Munch Media LLC
|
# Copyright (c) 2026 Rug Munch Media LLC
|
||||||
#
|
#
|
||||||
# Part of Pry — https://git.rugmunch.io/RugMunchMedia/pryscraper
|
# Part of Pry - https://git.rugmunch.io/RugMunchMedia/pryscraper
|
||||||
# Licensed under MIT. See LICENSE.
|
# Licensed under MIT. See LICENSE.
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
from collections.abc import Generator
|
import sqlite3
|
||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
from typing import Any
|
from pathlib import Path
|
||||||
|
from typing import Any, Iterator
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
# Try to import SQLAlchemy
|
|
||||||
try:
|
try:
|
||||||
from sqlalchemy import (
|
from sqlalchemy import (
|
||||||
JSON,
|
JSON,
|
||||||
|
|
@ -24,150 +73,601 @@ try:
|
||||||
Column,
|
Column,
|
||||||
DateTime,
|
DateTime,
|
||||||
Float,
|
Float,
|
||||||
|
ForeignKey,
|
||||||
|
Index,
|
||||||
Integer,
|
Integer,
|
||||||
String,
|
String,
|
||||||
Text,
|
Text,
|
||||||
create_engine,
|
create_engine,
|
||||||
|
event,
|
||||||
)
|
)
|
||||||
from sqlalchemy.dialects.postgresql import JSONB
|
from sqlalchemy.engine import Engine
|
||||||
from sqlalchemy.orm import Session, declarative_base, sessionmaker
|
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
|
||||||
_has_sqlalchemy = True
|
|
||||||
Base = declarative_base()
|
_HAS_SA = True
|
||||||
except ImportError:
|
except ImportError: # pragma: no cover
|
||||||
_has_sqlalchemy = False
|
_HAS_SA = False
|
||||||
Base = None
|
|
||||||
|
from paths import PRY_DATA_DIR # noqa: E402
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
DEFAULT_DB_PATH = PRY_DATA_DIR / "pry.db"
|
||||||
|
|
||||||
|
# Backward-compat aliases for existing code/tests that imported these names
|
||||||
|
_has_sqlalchemy = _HAS_SA
|
||||||
|
|
||||||
|
|
||||||
if _has_sqlalchemy:
|
def get_db() -> "Engine":
|
||||||
class ApiKey(Base):
|
"""Backward-compat alias for get_engine()."""
|
||||||
__tablename__ = "api_keys"
|
return get_engine()
|
||||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
|
||||||
key_hash = Column(String(64), unique=True, nullable=False, index=True)
|
|
||||||
user_id = Column(String(32), nullable=False, index=True)
|
|
||||||
name = Column(String(100), default="default")
|
|
||||||
rate_limit_rpm = Column(Integer, default=60)
|
|
||||||
created_at = Column(DateTime, default=lambda: datetime.now(UTC))
|
|
||||||
last_used = Column(DateTime, nullable=True)
|
|
||||||
use_count = Column(Integer, default=0)
|
|
||||||
active = Column(Boolean, default=True)
|
|
||||||
|
|
||||||
class User(Base):
|
DATABASE_URL_ENV = "PRY_DATABASE_URL"
|
||||||
__tablename__ = "users"
|
DEFAULT_SQLITE_URL = f"sqlite:///{DEFAULT_DB_PATH}"
|
||||||
id = Column(String(32), primary_key=True)
|
|
||||||
email = Column(String(255), unique=True, nullable=False, index=True)
|
|
||||||
password_hash = Column(String(128), nullable=False)
|
|
||||||
salt = Column(String(64), nullable=False)
|
|
||||||
role = Column(String(20), default="user")
|
|
||||||
created_at = Column(DateTime, default=lambda: datetime.now(UTC))
|
|
||||||
active = Column(Boolean, default=True)
|
|
||||||
metadata_json = Column("metadata", JSON, default=dict)
|
|
||||||
|
|
||||||
class UsageRecord(Base):
|
|
||||||
__tablename__ = "usage_records"
|
|
||||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
|
||||||
api_key_id = Column(Integer, index=True)
|
|
||||||
operation = Column(String(50), nullable=False, index=True)
|
|
||||||
quantity = Column(Float, default=1.0)
|
|
||||||
cost_usd = Column(Float, default=0.0)
|
|
||||||
timestamp = Column(DateTime, default=lambda: datetime.now(UTC), index=True)
|
|
||||||
metadata_json = Column("metadata", JSON, default=dict)
|
|
||||||
|
|
||||||
class QualityCheckRecord(Base):
|
def _resolve_database_url() -> str:
|
||||||
|
"""Resolve the database URL. Default: SQLite at PRY_DATA_DIR/pry.db."""
|
||||||
|
url = os.getenv(DATABASE_URL_ENV)
|
||||||
|
if url:
|
||||||
|
return url
|
||||||
|
return DEFAULT_SQLITE_URL
|
||||||
|
|
||||||
|
|
||||||
|
_engine: "Engine | None" = None
|
||||||
|
_SessionLocal: "sessionmaker[Session] | None" = None
|
||||||
|
|
||||||
|
|
||||||
|
def get_engine() -> "Engine":
|
||||||
|
"""Get the SQLAlchemy engine. Lazily creates one on first call."""
|
||||||
|
global _engine, _SessionLocal
|
||||||
|
if not _HAS_SA:
|
||||||
|
raise RuntimeError("sqlalchemy is not installed; pip install 'sqlalchemy>=2.0'")
|
||||||
|
if _engine is None:
|
||||||
|
url = _resolve_database_url()
|
||||||
|
connect_args: dict[str, Any] = {}
|
||||||
|
if url.startswith("sqlite"):
|
||||||
|
connect_args["check_same_thread"] = False
|
||||||
|
_engine = create_engine(url, connect_args=connect_args, future=True, echo=False)
|
||||||
|
# Enable foreign keys for SQLite (off by default)
|
||||||
|
if url.startswith("sqlite"):
|
||||||
|
@event.listens_for(_engine, "connect")
|
||||||
|
def _enable_sqlite_fk(dbapi_conn: sqlite3.Connection, _conn_record: Any) -> None:
|
||||||
|
cur = dbapi_conn.cursor()
|
||||||
|
cur.execute("PRAGMA foreign_keys=ON")
|
||||||
|
cur.close()
|
||||||
|
_SessionLocal = sessionmaker(bind=_engine, autoflush=False, expire_on_commit=False)
|
||||||
|
# Auto-create tables on first import (idempotent)
|
||||||
|
Base.metadata.create_all(_engine)
|
||||||
|
logger.info("db_engine_initialized", extra={"url": url.split("@")[-1] if "@" in url else url})
|
||||||
|
return _engine
|
||||||
|
|
||||||
|
|
||||||
|
def get_session() -> "Session":
|
||||||
|
"""Get a new Session. Caller is responsible for closing it.
|
||||||
|
|
||||||
|
Use `with db_session() as s:` or the `session_scope()` context manager
|
||||||
|
for automatic cleanup and commit/rollback handling.
|
||||||
|
"""
|
||||||
|
if _SessionLocal is None:
|
||||||
|
get_engine()
|
||||||
|
assert _SessionLocal is not None
|
||||||
|
return _SessionLocal()
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def session_scope() -> Iterator["Session"]:
|
||||||
|
"""Context manager: yields a Session, commits on success, rolls back on error.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
with session_scope() as s:
|
||||||
|
s.add(QualityCheck(...))
|
||||||
|
# commit happens automatically when the with-block exits cleanly
|
||||||
|
"""
|
||||||
|
s = get_session()
|
||||||
|
try:
|
||||||
|
yield s
|
||||||
|
s.commit()
|
||||||
|
except Exception:
|
||||||
|
s.rollback()
|
||||||
|
raise
|
||||||
|
finally:
|
||||||
|
s.close()
|
||||||
|
|
||||||
|
|
||||||
|
# ── Model base ────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
if _HAS_SA:
|
||||||
|
class Base(DeclarativeBase):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _now() -> datetime:
|
||||||
|
return datetime.now(UTC)
|
||||||
|
|
||||||
|
# ── Quality ──
|
||||||
|
class QualityCheck(Base):
|
||||||
__tablename__ = "quality_checks"
|
__tablename__ = "quality_checks"
|
||||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
id = Column(Integer, primary_key=True)
|
||||||
url_hash = Column(String(16), index=True)
|
extraction_id = Column(String(64), index=True, nullable=False)
|
||||||
url = Column(String(2048), nullable=False)
|
url = Column(Text, nullable=False)
|
||||||
|
completeness = Column(Float, default=0.0)
|
||||||
|
accuracy = Column(Float, default=0.0)
|
||||||
|
freshness = Column(Float, default=0.0)
|
||||||
|
overall_score = Column(Float, default=0.0)
|
||||||
|
risk_level = Column(String(16), default="unknown")
|
||||||
|
issues = Column(JSON, default=list)
|
||||||
|
checked_at = Column(DateTime, default=_now, index=True)
|
||||||
data = Column(JSON, default=dict)
|
data = Column(JSON, default=dict)
|
||||||
quality_score = Column(Float, default=0.0)
|
|
||||||
anomaly_count = Column(Integer, default=0)
|
|
||||||
checked_at = Column(DateTime, default=lambda: datetime.now(UTC))
|
|
||||||
|
|
||||||
class ReviewRequest(Base):
|
# ── Reviews (human review queue) ──
|
||||||
__tablename__ = "review_requests"
|
class ReviewItem(Base):
|
||||||
id = Column(String(12), primary_key=True)
|
__tablename__ = "review_items"
|
||||||
user_id = Column(String(32), index=True)
|
id = Column(Integer, primary_key=True)
|
||||||
status = Column(String(20), default="pending", index=True)
|
review_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
data = Column(JSON, default=dict)
|
url = Column(Text, nullable=False)
|
||||||
confidence_score = Column(Float, default=0.0)
|
fields = Column(JSON, default=dict)
|
||||||
flagged_fields = Column(JSON, default=list)
|
status = Column(String(16), default="pending", index=True) # pending|approved|rejected
|
||||||
created_at = Column(DateTime, default=lambda: datetime.now(UTC))
|
reviewer = Column(String(128), default="")
|
||||||
reviewed_at = Column(DateTime, nullable=True)
|
notes = Column(Text, default="")
|
||||||
reviewed_by = Column(String(100), nullable=True)
|
submitted_at = Column(DateTime, default=_now, index=True)
|
||||||
review_notes = Column(Text, default="")
|
resolved_at = Column(DateTime, nullable=True)
|
||||||
|
|
||||||
class MonitorRecord(Base):
|
# ── Intel (competitive intelligence snapshots) ──
|
||||||
|
class IntelSnapshot(Base):
|
||||||
|
__tablename__ = "intel_snapshots"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
competitor_id = Column(String(64), index=True, nullable=False)
|
||||||
|
competitor_name = Column(String(256), nullable=False)
|
||||||
|
url = Column(Text, nullable=False)
|
||||||
|
fields = Column(JSON, default=dict)
|
||||||
|
ts = Column(DateTime, default=_now, index=True)
|
||||||
|
Index("ix_intel_competitor_ts", competitor_id, ts.desc())
|
||||||
|
|
||||||
|
# ── Costing ──
|
||||||
|
class CostingEntry(Base):
|
||||||
|
__tablename__ = "costing_entries"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
operation = Column(String(64), index=True, nullable=False)
|
||||||
|
cost_usd = Column(Float, default=0.0)
|
||||||
|
units = Column(Float, default=1.0)
|
||||||
|
provider = Column(String(64), default="")
|
||||||
|
metadata_json = Column("metadata", JSON, default=dict) # 'metadata' is reserved
|
||||||
|
ts = Column(DateTime, default=_now, index=True)
|
||||||
|
|
||||||
|
# ── Freshness ──
|
||||||
|
class FreshnessSnapshot(Base):
|
||||||
|
__tablename__ = "freshness_snapshots"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
url = Column(Text, index=True, nullable=False)
|
||||||
|
content_hash = Column(String(64), default="")
|
||||||
|
fingerprint = Column(String(128), default="")
|
||||||
|
checked_at = Column(DateTime, default=_now, index=True)
|
||||||
|
|
||||||
|
# ── Structure (page structure monitor) ──
|
||||||
|
class StructureSnapshot(Base):
|
||||||
|
__tablename__ = "structure_snapshots"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
url = Column(Text, index=True, nullable=False)
|
||||||
|
structure_hash = Column(String(64), default="")
|
||||||
|
dom_summary = Column(JSON, default=dict)
|
||||||
|
checked_at = Column(DateTime, default=_now, index=True)
|
||||||
|
|
||||||
|
# ── SEO ──
|
||||||
|
class SeoSnapshot(Base):
|
||||||
|
__tablename__ = "seo_snapshots"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
url = Column(Text, index=True, nullable=False)
|
||||||
|
title = Column(Text, default="")
|
||||||
|
meta_description = Column(Text, default="")
|
||||||
|
h1 = Column(Text, default="")
|
||||||
|
keyword_count = Column(Integer, default=0)
|
||||||
|
readability_score = Column(Float, default=0.0)
|
||||||
|
full_data = Column(JSON, default=dict)
|
||||||
|
checked_at = Column(DateTime, default=_now, index=True)
|
||||||
|
|
||||||
|
# ── Monitors ──
|
||||||
|
class Monitor(Base):
|
||||||
__tablename__ = "monitors"
|
__tablename__ = "monitors"
|
||||||
id = Column(String(12), primary_key=True)
|
id = Column(Integer, primary_key=True)
|
||||||
user_id = Column(String(32), index=True)
|
monitor_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
name = Column(String(200), nullable=False)
|
url = Column(Text, nullable=False)
|
||||||
target_url = Column(String(2048), nullable=False)
|
monitor_name = Column(String(256), default="") # was 'name' in JSON
|
||||||
schedule_cron = Column(String(50), default="0 */6 * * *")
|
schedule_cron = Column(String(64), default="")
|
||||||
goal = Column(Text, default="")
|
check_type = Column(String(32), default="content")
|
||||||
webhook_url = Column(String(2048), default="")
|
active = Column(Boolean, default=True, index=True)
|
||||||
status = Column(String(20), default="active")
|
webhook_url = Column(Text, default="")
|
||||||
created_at = Column(DateTime, default=lambda: datetime.now(UTC))
|
last_check_at = Column(DateTime, nullable=True)
|
||||||
last_run_at = Column(DateTime, nullable=True)
|
created_at = Column(DateTime, default=_now)
|
||||||
total_checks = Column(Integer, default=0)
|
|
||||||
total_changes = Column(Integer, default=0)
|
class MonitorRun(Base):
|
||||||
|
__tablename__ = "monitor_runs"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
monitor_id = Column(String(64), ForeignKey("monitors.monitor_id"), index=True, nullable=False)
|
||||||
|
ts = Column(DateTime, default=_now, index=True)
|
||||||
|
status = Column(String(16), default="success") # success|changed|error
|
||||||
|
diff = Column(JSON, default=dict)
|
||||||
|
error = Column(Text, default="")
|
||||||
|
|
||||||
|
# ── Accounts (account pool for signup automation) ──
|
||||||
|
class Account(Base):
|
||||||
|
__tablename__ = "accounts"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
account_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
service = Column(String(64), index=True, nullable=False)
|
||||||
|
username = Column(String(256), nullable=False)
|
||||||
|
email = Column(String(256), default="")
|
||||||
|
status = Column(String(16), default="active") # active|cooldown|disabled
|
||||||
|
used_count = Column(Integer, default=0)
|
||||||
|
last_used_at = Column(DateTime, nullable=True)
|
||||||
|
created_at = Column(DateTime, default=_now)
|
||||||
|
|
||||||
|
# ── Browser sessions ──
|
||||||
|
class BrowserSession(Base):
|
||||||
|
__tablename__ = "browser_sessions"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
session_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
cookies = Column(JSON, default=list)
|
||||||
|
local_storage = Column(JSON, default=dict)
|
||||||
|
user_agent = Column(String(512), default="")
|
||||||
|
proxy = Column(String(256), default="")
|
||||||
|
created_at = Column(DateTime, default=_now)
|
||||||
|
expires_at = Column(DateTime, nullable=True)
|
||||||
|
|
||||||
|
# ── Reports ──
|
||||||
|
class Report(Base):
|
||||||
|
__tablename__ = "reports"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
report_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
report_name = Column(String(256), default="") # was 'name' in JSON
|
||||||
|
report_type = Column(String(64), default="html") # html|pdf|csv
|
||||||
|
url = Column(Text, default="")
|
||||||
|
data = Column(JSON, default=dict)
|
||||||
|
generated_at = Column(DateTime, default=_now, index=True)
|
||||||
|
|
||||||
|
# ── Training datasets ──
|
||||||
|
class TrainingDataset(Base):
|
||||||
|
__tablename__ = "training_datasets"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
dataset_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
dataset_name = Column(String(256), default="") # was 'name' in JSON
|
||||||
|
license_class = Column(String(64), default="unknown")
|
||||||
|
pii_redacted = Column(Boolean, default=False)
|
||||||
|
record_count = Column(Integer, default=0)
|
||||||
|
file_path = Column(Text, default="")
|
||||||
|
created_at = Column(DateTime, default=_now)
|
||||||
|
|
||||||
|
# ── Pipelines ──
|
||||||
|
class Pipeline(Base):
|
||||||
|
__tablename__ = "pipelines"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
pipeline_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
pipeline_name = Column(String(256), default="") # was 'name' in JSON
|
||||||
|
steps = Column(JSON, default=list)
|
||||||
|
schedule_cron = Column(String(64), default="")
|
||||||
|
active = Column(Boolean, default=True)
|
||||||
|
created_at = Column(DateTime, default=_now)
|
||||||
|
|
||||||
|
class PipelineRun(Base):
|
||||||
|
__tablename__ = "pipeline_runs"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
pipeline_id = Column(String(64), ForeignKey("pipelines.pipeline_id"), index=True, nullable=False)
|
||||||
|
ts = Column(DateTime, default=_now, index=True)
|
||||||
|
status = Column(String(16), default="success")
|
||||||
|
result = Column(JSON, default=dict)
|
||||||
|
|
||||||
|
# ── GDPR ──
|
||||||
|
class GdprRequest(Base):
|
||||||
|
__tablename__ = "gdpr_requests"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
request_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
request_type = Column(String(32), nullable=False) # access|delete|portability
|
||||||
|
subject_id = Column(String(256), nullable=False, index=True)
|
||||||
|
status = Column(String(16), default="pending")
|
||||||
|
requested_at = Column(DateTime, default=_now, index=True)
|
||||||
|
completed_at = Column(DateTime, nullable=True)
|
||||||
|
details = Column(JSON, default=dict)
|
||||||
|
|
||||||
|
# ── Agency (white-label) ──
|
||||||
|
class Agency(Base):
|
||||||
|
__tablename__ = "agencies"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
agency_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
agency_name = Column(String(256), nullable=False) # was 'name' in JSON
|
||||||
|
owner_email = Column(String(256), nullable=False)
|
||||||
|
custom_domain = Column(String(256), default="")
|
||||||
|
brand_color = Column(String(16), default="#f59e0b")
|
||||||
|
logo_url = Column(Text, default="")
|
||||||
|
quota = Column(Integer, default=10000)
|
||||||
|
active = Column(Boolean, default=True)
|
||||||
|
created_at = Column(DateTime, default=_now)
|
||||||
|
|
||||||
class AgencyClient(Base):
|
class AgencyClient(Base):
|
||||||
__tablename__ = "agency_clients"
|
__tablename__ = "agency_clients"
|
||||||
id = Column(String(12), primary_key=True)
|
id = Column(Integer, primary_key=True)
|
||||||
agency_id = Column(String(32), index=True)
|
client_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
name = Column(String(200), nullable=False)
|
agency_id = Column(String(64), ForeignKey("agencies.agency_id"), index=True, nullable=False)
|
||||||
email = Column(String(255), nullable=False)
|
name = Column(String(256), nullable=False)
|
||||||
api_key_hash = Column(String(64), unique=True)
|
email = Column(String(256), default="")
|
||||||
monthly_quota = Column(Integer, default=10000)
|
quota_used = Column(Integer, default=0)
|
||||||
usage_this_month = Column(Integer, default=0)
|
quota_limit = Column(Integer, default=1000)
|
||||||
status = Column(String(20), default="active")
|
active = Column(Boolean, default=True)
|
||||||
created_at = Column(DateTime, default=lambda: datetime.now(UTC))
|
created_at = Column(DateTime, default=_now)
|
||||||
|
|
||||||
|
# ── Referrals (click/conversion tracking) ──
|
||||||
|
class ReferralClick(Base):
|
||||||
|
__tablename__ = "referral_clicks"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
click_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
provider = Column(String(64), index=True, nullable=False)
|
||||||
|
url = Column(Text, nullable=False)
|
||||||
|
source = Column(String(64), default="")
|
||||||
|
user_id = Column(String(128), default="")
|
||||||
|
ts = Column(DateTime, default=_now, index=True)
|
||||||
|
converted = Column(Boolean, default=False)
|
||||||
|
conversion_value_usd = Column(Float, default=0.0)
|
||||||
|
|
||||||
|
# ── Actors (Apify-style marketplace) ──
|
||||||
|
class Actor(Base):
|
||||||
|
__tablename__ = "actors"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
actor_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
name = Column(String(256), nullable=False)
|
||||||
|
description = Column(Text, default="")
|
||||||
|
template_id = Column(String(64), default="")
|
||||||
|
code = Column(Text, default="")
|
||||||
|
price_per_run = Column(Float, default=0.0)
|
||||||
|
visibility = Column(String(16), default="private") # private|public|unlisted
|
||||||
|
schedule_cron = Column(String(64), default="")
|
||||||
|
tags = Column(JSON, default=list)
|
||||||
|
author = Column(String(128), default="")
|
||||||
|
run_count = Column(Integer, default=0)
|
||||||
|
revenue_usd = Column(Float, default=0.0)
|
||||||
|
created_at = Column(DateTime, default=_now)
|
||||||
|
|
||||||
|
class ActorRun(Base):
|
||||||
|
__tablename__ = "actor_runs"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
actor_id = Column(String(64), ForeignKey("actors.actor_id"), index=True, nullable=False)
|
||||||
|
ts = Column(DateTime, default=_now, index=True)
|
||||||
|
status = Column(String(16), default="success")
|
||||||
|
result = Column(JSON, default=dict)
|
||||||
|
cost_usd = Column(Float, default=0.0)
|
||||||
|
|
||||||
|
# ── LLM usage tracking ──
|
||||||
|
class LlmUsage(Base):
|
||||||
|
__tablename__ = "llm_usage"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
provider = Column(String(64), index=True, nullable=False)
|
||||||
|
model = Column(String(128), default="")
|
||||||
|
input_tokens = Column(Integer, default=0)
|
||||||
|
output_tokens = Column(Integer, default=0)
|
||||||
|
cost_usd = Column(Float, default=0.0)
|
||||||
|
operation = Column(String(64), default="")
|
||||||
|
ts = Column(DateTime, default=_now, index=True)
|
||||||
|
|
||||||
|
# ── Webhooks ──
|
||||||
|
class Webhook(Base):
|
||||||
|
__tablename__ = "webhooks"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
webhook_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
url = Column(Text, nullable=False)
|
||||||
|
events = Column(JSON, default=list)
|
||||||
|
secret = Column(String(128), default="")
|
||||||
|
active = Column(Boolean, default=True)
|
||||||
|
created_at = Column(DateTime, default=_now)
|
||||||
|
|
||||||
|
# ── x402 receipts ──
|
||||||
|
class X402Receipt(Base):
|
||||||
|
__tablename__ = "x402_receipts"
|
||||||
|
id = Column(Integer, primary_key=True)
|
||||||
|
receipt_id = Column(String(64), unique=True, index=True, nullable=False)
|
||||||
|
tx_hash = Column(String(128), index=True, default="")
|
||||||
|
network = Column(String(32), default="base")
|
||||||
|
asset = Column(String(16), default="USDC")
|
||||||
|
amount_usd = Column(Float, default=0.0)
|
||||||
|
payer = Column(String(128), default="")
|
||||||
|
payee = Column(String(128), default="")
|
||||||
|
operation = Column(String(64), default="")
|
||||||
|
ts = Column(DateTime, default=_now, index=True)
|
||||||
|
facilitator = Column(String(64), default="x402.org")
|
||||||
|
|
||||||
|
else: # pragma: no cover
|
||||||
|
Base = None # type: ignore
|
||||||
|
|
||||||
|
|
||||||
class Database:
|
# ── Importer: JSON store -> SQL ───────────────────────────────────
|
||||||
"""SQLAlchemy database wrapper. Supports SQLite (default) and PostgreSQL."""
|
|
||||||
|
|
||||||
def __init__(self, url: str = ""):
|
def _iter_jsonl(path: Path) -> Iterator[dict[str, Any]]:
|
||||||
if not _has_sqlalchemy:
|
if not path.exists():
|
||||||
logger.warning("sqlalchemy_not_available")
|
return
|
||||||
return
|
try:
|
||||||
if not url:
|
for line in path.read_text(encoding="utf-8").splitlines():
|
||||||
url = os.getenv("PRY_DATABASE_URL", "sqlite:///~/.pry/data.db")
|
if not line.strip():
|
||||||
# Ensure directory exists
|
continue
|
||||||
if url.startswith("sqlite:///"):
|
try:
|
||||||
db_path = url[10:]
|
yield json.loads(line)
|
||||||
os.makedirs(os.path.dirname(os.path.expanduser(db_path)), exist_ok=True)
|
except json.JSONDecodeError:
|
||||||
url = f"sqlite:///{os.path.expanduser(db_path)}"
|
continue
|
||||||
self.engine = create_engine(url, echo=False, pool_pre_ping=True)
|
except OSError as e:
|
||||||
self.Session = sessionmaker(bind=self.engine)
|
logger.warning("jsonl_read_failed", extra={"path": str(path), "error": str(e)[:80]})
|
||||||
# Create tables
|
|
||||||
Base.metadata.create_all(self.engine)
|
|
||||||
logger.info("database_initialized", extra={"url": url.split("@")[-1]})
|
|
||||||
|
|
||||||
@contextmanager
|
|
||||||
def session(self) -> Generator[Any, None, None]:
|
def _iter_json_files(path: Path) -> Iterator[dict[str, Any]]:
|
||||||
if not _has_sqlalchemy:
|
if not path.exists():
|
||||||
yield None
|
return
|
||||||
return
|
if path.is_file():
|
||||||
s = self.Session()
|
|
||||||
try:
|
try:
|
||||||
yield s
|
data = json.loads(path.read_text(encoding="utf-8"))
|
||||||
s.commit()
|
if isinstance(data, list):
|
||||||
except Exception: # noqa: BLE001
|
for d in data:
|
||||||
s.rollback()
|
if isinstance(d, dict):
|
||||||
raise
|
yield d
|
||||||
finally:
|
elif isinstance(data, dict):
|
||||||
s.close()
|
yield data
|
||||||
|
except (json.JSONDecodeError, OSError):
|
||||||
def is_available(self) -> bool:
|
return
|
||||||
return _has_sqlalchemy
|
return
|
||||||
|
for f in path.glob("*.json"):
|
||||||
|
yield from _iter_json_files(f)
|
||||||
|
|
||||||
|
|
||||||
_db: Database | None = None
|
def import_json_stores(data_dir: Path | None = None) -> dict[str, int]:
|
||||||
|
"""One-shot importer: read all JSON stores and write to the SQL tables.
|
||||||
|
|
||||||
|
Returns a dict of {store_name: row_count_imported}.
|
||||||
|
|
||||||
|
Idempotent: re-running with the same data is safe (uses upsert-style
|
||||||
|
pattern via the unique constraints on _id columns).
|
||||||
|
"""
|
||||||
|
data_dir = data_dir or PRY_DATA_DIR
|
||||||
|
counts: dict[str, int] = {}
|
||||||
|
|
||||||
|
with session_scope() as s:
|
||||||
|
# Quality
|
||||||
|
n = 0
|
||||||
|
for rec in _iter_jsonl(data_dir / "quality" / "checks.jsonl"):
|
||||||
|
s.add(QualityCheck(
|
||||||
|
extraction_id=rec.get("extraction_id", rec.get("id", ""))[:64],
|
||||||
|
url=rec.get("url", ""),
|
||||||
|
completeness=rec.get("completeness", 0.0),
|
||||||
|
accuracy=rec.get("accuracy", 0.0),
|
||||||
|
freshness=rec.get("freshness", 0.0),
|
||||||
|
overall_score=rec.get("overall_score", 0.0),
|
||||||
|
risk_level=rec.get("risk_level", "unknown"),
|
||||||
|
issues=rec.get("issues", []),
|
||||||
|
data=rec,
|
||||||
|
))
|
||||||
|
n += 1
|
||||||
|
counts["quality"] = n
|
||||||
|
|
||||||
|
# Intel (JSONL)
|
||||||
|
n = 0
|
||||||
|
for rec in _iter_jsonl(data_dir / "intel" / "snapshots.jsonl"):
|
||||||
|
s.add(IntelSnapshot(
|
||||||
|
competitor_id=rec.get("competitor_id", "")[:64],
|
||||||
|
competitor_name=rec.get("competitor_name", "")[:256],
|
||||||
|
url=rec.get("url", ""),
|
||||||
|
fields=rec.get("fields", {}),
|
||||||
|
))
|
||||||
|
n += 1
|
||||||
|
counts["intel"] = n
|
||||||
|
|
||||||
|
# Monitors (JSON files)
|
||||||
|
n = 0
|
||||||
|
for rec in _iter_json_files(data_dir / "monitors"):
|
||||||
|
mid = rec.get("monitor_id") or rec.get("id", "")
|
||||||
|
s.add(Monitor(
|
||||||
|
monitor_id=mid[:64],
|
||||||
|
url=rec.get("url", ""),
|
||||||
|
monitor_name=rec.get("name", ""), # was 'name' in JSON
|
||||||
|
schedule_cron=rec.get("schedule_cron", ""),
|
||||||
|
check_type=rec.get("check_type", "content"),
|
||||||
|
active=rec.get("active", True),
|
||||||
|
webhook_url=rec.get("webhook_url", ""),
|
||||||
|
))
|
||||||
|
n += 1
|
||||||
|
counts["monitors"] = n
|
||||||
|
|
||||||
|
# Accounts (JSON files)
|
||||||
|
n = 0
|
||||||
|
for rec in _iter_json_files(data_dir / "accounts"):
|
||||||
|
aid = rec.get("account_id") or rec.get("id", "")
|
||||||
|
s.add(Account(
|
||||||
|
account_id=aid[:64],
|
||||||
|
service=rec.get("service", "unknown")[:64],
|
||||||
|
username=rec.get("username", ""),
|
||||||
|
email=rec.get("email", ""),
|
||||||
|
status=rec.get("status", "active"),
|
||||||
|
used_count=rec.get("used_count", 0),
|
||||||
|
))
|
||||||
|
n += 1
|
||||||
|
counts["accounts"] = n
|
||||||
|
|
||||||
|
# Actors (JSON files)
|
||||||
|
n = 0
|
||||||
|
for rec in _iter_json_files(data_dir / "actors"):
|
||||||
|
aid = rec.get("actor_id") or rec.get("id", "")
|
||||||
|
s.add(Actor(
|
||||||
|
actor_id=aid[:64],
|
||||||
|
name=rec.get("name", ""),
|
||||||
|
description=rec.get("description", ""),
|
||||||
|
template_id=rec.get("template_id", ""),
|
||||||
|
code=rec.get("code", ""),
|
||||||
|
price_per_run=rec.get("price_per_run", 0.0),
|
||||||
|
visibility=rec.get("visibility", "private"),
|
||||||
|
schedule_cron=rec.get("schedule_cron", ""),
|
||||||
|
tags=rec.get("tags", []),
|
||||||
|
author=rec.get("author", ""),
|
||||||
|
run_count=rec.get("run_count", 0),
|
||||||
|
revenue_usd=rec.get("revenue_usd", 0.0),
|
||||||
|
))
|
||||||
|
n += 1
|
||||||
|
counts["actors"] = n
|
||||||
|
|
||||||
|
# Pipelines (JSON files)
|
||||||
|
n = 0
|
||||||
|
for rec in _iter_json_files(data_dir / "pipelines"):
|
||||||
|
pid = rec.get("pipeline_id") or rec.get("id", "")
|
||||||
|
s.add(Pipeline(
|
||||||
|
pipeline_id=pid[:64],
|
||||||
|
pipeline_name=rec.get("name", ""),
|
||||||
|
steps=rec.get("steps", []),
|
||||||
|
schedule_cron=rec.get("schedule_cron", ""),
|
||||||
|
active=rec.get("active", True),
|
||||||
|
))
|
||||||
|
n += 1
|
||||||
|
counts["pipelines"] = n
|
||||||
|
|
||||||
|
# Sessions (JSON files)
|
||||||
|
n = 0
|
||||||
|
for rec in _iter_json_files(data_dir / "sessions"):
|
||||||
|
sid = rec.get("session_id") or rec.get("id", "")
|
||||||
|
s.add(BrowserSession(
|
||||||
|
session_id=sid[:64],
|
||||||
|
cookies=rec.get("cookies", []),
|
||||||
|
local_storage=rec.get("local_storage", {}),
|
||||||
|
user_agent=rec.get("user_agent", ""),
|
||||||
|
proxy=rec.get("proxy", ""),
|
||||||
|
))
|
||||||
|
n += 1
|
||||||
|
counts["sessions"] = n
|
||||||
|
|
||||||
|
# Agency + clients (JSON files)
|
||||||
|
n = 0
|
||||||
|
for rec in _iter_json_files(data_dir / "agency"):
|
||||||
|
aid = rec.get("agency_id") or rec.get("id", "")
|
||||||
|
s.add(Agency(
|
||||||
|
agency_id=aid[:64],
|
||||||
|
agency_name=rec.get("name", ""),
|
||||||
|
owner_email=rec.get("owner_email", ""),
|
||||||
|
custom_domain=rec.get("custom_domain", ""),
|
||||||
|
brand_color=rec.get("brand_color", "#f59e0b"),
|
||||||
|
logo_url=rec.get("logo_url", ""),
|
||||||
|
quota=rec.get("quota", 10000),
|
||||||
|
))
|
||||||
|
n += 1
|
||||||
|
counts["agency"] = n
|
||||||
|
|
||||||
|
return counts
|
||||||
|
|
||||||
|
|
||||||
def get_db() -> Database | None:
|
def db_health() -> dict[str, Any]:
|
||||||
"""Get or create the global database instance."""
|
"""Return a quick health dict for /health or /status endpoints."""
|
||||||
global _db
|
if not _HAS_SA:
|
||||||
if _db is None and _has_sqlalchemy:
|
return {"available": False, "reason": "sqlalchemy not installed"}
|
||||||
_db = Database()
|
try:
|
||||||
return _db
|
eng = get_engine()
|
||||||
|
with eng.connect() as conn:
|
||||||
|
# SQLite-specific pragma; no-op for other backends
|
||||||
|
try:
|
||||||
|
result = conn.exec_driver_sql("SELECT sqlite_version()").fetchone()
|
||||||
|
sqlite_version = result[0] if result else None
|
||||||
|
except Exception:
|
||||||
|
sqlite_version = None
|
||||||
|
return {
|
||||||
|
"available": True,
|
||||||
|
"url": _resolve_database_url().split("@")[-1] if "@" in _resolve_database_url() else _resolve_database_url(),
|
||||||
|
"sqlite_version": sqlite_version,
|
||||||
|
}
|
||||||
|
except Exception as e:
|
||||||
|
return {"available": False, "error": str(e)[:200]}
|
||||||
|
|
|
||||||
|
|
@ -51,6 +51,8 @@ dependencies = [
|
||||||
"anyio>=4.0.0",
|
"anyio>=4.0.0",
|
||||||
"croniter>=2.0.0",
|
"croniter>=2.0.0",
|
||||||
"structlog>=24.0.0",
|
"structlog>=24.0.0",
|
||||||
|
"sqlalchemy>=2.0.0",
|
||||||
|
"aiosqlite>=0.19.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.optional-dependencies]
|
[project.optional-dependencies]
|
||||||
|
|
|
||||||
167
tests/test_db.py
Normal file
167
tests/test_db.py
Normal file
|
|
@ -0,0 +1,167 @@
|
||||||
|
"""Tests for the db module (SQLAlchemy foundation).
|
||||||
|
|
||||||
|
Uses a temp directory for PRY_DATA_DIR so we don't pollute the user's
|
||||||
|
real data dir. Tests verify schema creation, basic CRUD, and the
|
||||||
|
JSON store importer.
|
||||||
|
"""
|
||||||
|
# SPDX-License-Identifier: MIT
|
||||||
|
# Copyright (c) 2026 Rug Munch Media LLC
|
||||||
|
#
|
||||||
|
# Part of Pry - https://git.rugmunch.io/RugMunchMedia/pryscraper
|
||||||
|
# Licensed under MIT. See LICENSE.
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import importlib
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def temp_data_dir(monkeypatch, tmp_path):
|
||||||
|
"""Set PRY_DATA_DIR to a temp dir and reload db to pick up the change."""
|
||||||
|
monkeypatch.setenv("PRY_DATA_DIR", str(tmp_path))
|
||||||
|
monkeypatch.setenv("PRY_DATABASE_URL", f"sqlite:///{tmp_path}/test.db")
|
||||||
|
# Force re-import of paths and db
|
||||||
|
import paths
|
||||||
|
import db
|
||||||
|
importlib.reload(paths)
|
||||||
|
importlib.reload(db)
|
||||||
|
# Reset the cached engine
|
||||||
|
db._engine = None
|
||||||
|
db._SessionLocal = None
|
||||||
|
yield tmp_path
|
||||||
|
# Cleanup
|
||||||
|
db._engine = None
|
||||||
|
db._SessionLocal = None
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_engine_creates_db_file(temp_data_dir):
|
||||||
|
import db
|
||||||
|
eng = db.get_engine()
|
||||||
|
assert eng is not None
|
||||||
|
# The DB file should exist now
|
||||||
|
assert (temp_data_dir / "test.db").exists()
|
||||||
|
|
||||||
|
|
||||||
|
def test_all_tables_created(temp_data_dir):
|
||||||
|
import db
|
||||||
|
from sqlalchemy import inspect
|
||||||
|
db.get_engine()
|
||||||
|
inspector = inspect(db.get_engine())
|
||||||
|
tables = inspector.get_table_names()
|
||||||
|
# Should have all 24 model tables
|
||||||
|
expected_tables = {
|
||||||
|
"quality_checks", "review_items", "intel_snapshots", "costing_entries",
|
||||||
|
"freshness_snapshots", "structure_snapshots", "seo_snapshots", "monitors",
|
||||||
|
"monitor_runs", "accounts", "browser_sessions", "reports", "training_datasets",
|
||||||
|
"pipelines", "pipeline_runs", "gdpr_requests", "agencies", "agency_clients",
|
||||||
|
"referral_clicks", "actors", "actor_runs", "llm_usage", "webhooks", "x402_receipts",
|
||||||
|
}
|
||||||
|
assert expected_tables.issubset(set(tables)), f"missing: {expected_tables - set(tables)}"
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_scope_commits_on_success(temp_data_dir):
|
||||||
|
import db
|
||||||
|
with db.session_scope() as s:
|
||||||
|
s.add(db.QualityCheck(
|
||||||
|
extraction_id="test-1",
|
||||||
|
url="https://example.com",
|
||||||
|
completeness=0.9,
|
||||||
|
))
|
||||||
|
# Verify it was committed
|
||||||
|
with db.session_scope() as s:
|
||||||
|
from sqlalchemy import select
|
||||||
|
result = s.execute(select(db.QualityCheck).where(db.QualityCheck.extraction_id == "test-1"))
|
||||||
|
row = result.scalar_one_or_none()
|
||||||
|
assert row is not None
|
||||||
|
assert row.url == "https://example.com"
|
||||||
|
assert row.completeness == 0.9
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_scope_rolls_back_on_error(temp_data_dir):
|
||||||
|
import db
|
||||||
|
try:
|
||||||
|
with db.session_scope() as s:
|
||||||
|
s.add(db.QualityCheck(extraction_id="test-2", url="https://example.com", completeness=0.5))
|
||||||
|
raise ValueError("intentional")
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
# Verify it was NOT committed
|
||||||
|
with db.session_scope() as s:
|
||||||
|
from sqlalchemy import select
|
||||||
|
result = s.execute(select(db.QualityCheck).where(db.QualityCheck.extraction_id == "test-2"))
|
||||||
|
assert result.scalar_one_or_none() is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_import_json_stores_reads_existing_data(temp_data_dir):
|
||||||
|
"""The importer should read existing JSON files and write to SQL tables."""
|
||||||
|
# Set up some JSON files
|
||||||
|
intel_dir = temp_data_dir / "intel"
|
||||||
|
intel_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
(intel_dir / "snapshots.jsonl").write_text( # use the canonical filename
|
||||||
|
json.dumps({
|
||||||
|
"competitor_id": "acme",
|
||||||
|
"competitor_name": "Acme Corp",
|
||||||
|
"url": "https://acme.com",
|
||||||
|
"fields": {"price": 99.99},
|
||||||
|
}) + "\n" +
|
||||||
|
json.dumps({
|
||||||
|
"competitor_id": "acme",
|
||||||
|
"competitor_name": "Acme Corp",
|
||||||
|
"url": "https://acme.com/pricing",
|
||||||
|
"fields": {"price": 89.99},
|
||||||
|
}) + "\n"
|
||||||
|
)
|
||||||
|
|
||||||
|
monitor_dir = temp_data_dir / "monitors"
|
||||||
|
monitor_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
(monitor_dir / "m1.json").write_text(json.dumps({
|
||||||
|
"monitor_id": "m1",
|
||||||
|
"url": "https://example.com",
|
||||||
|
"name": "Example Monitor",
|
||||||
|
"schedule_cron": "0 * * * *",
|
||||||
|
"active": True,
|
||||||
|
}))
|
||||||
|
|
||||||
|
import db
|
||||||
|
counts = db.import_json_stores(data_dir=temp_data_dir)
|
||||||
|
assert counts.get("intel") == 2
|
||||||
|
assert counts.get("monitors") == 1
|
||||||
|
assert counts.get("actors") == 0 # empty dir
|
||||||
|
|
||||||
|
# Verify in SQL
|
||||||
|
with db.session_scope() as s:
|
||||||
|
from sqlalchemy import select
|
||||||
|
intel_rows = s.execute(select(db.IntelSnapshot).where(db.IntelSnapshot.competitor_id == "acme")).scalars().all()
|
||||||
|
assert len(intel_rows) == 2
|
||||||
|
monitor_row = s.execute(select(db.Monitor).where(db.Monitor.monitor_id == "m1")).scalar_one_or_none()
|
||||||
|
assert monitor_row is not None
|
||||||
|
# 'name' from JSON was renamed to 'monitor_name' in SQL
|
||||||
|
assert monitor_row.monitor_name == "Example Monitor"
|
||||||
|
|
||||||
|
|
||||||
|
def test_db_health_returns_dict(temp_data_dir):
|
||||||
|
import db
|
||||||
|
health = db.db_health()
|
||||||
|
assert isinstance(health, dict)
|
||||||
|
assert "available" in health
|
||||||
|
assert health["available"] is True
|
||||||
|
assert "url" in health
|
||||||
|
|
||||||
|
|
||||||
|
def test_models_have_unique_constraints(temp_data_dir):
|
||||||
|
"""Tables that map to JSON stores with _id fields should have unique indexes
|
||||||
|
so that re-importing the same data is idempotent (or at least detectable)."""
|
||||||
|
import db
|
||||||
|
from sqlalchemy import inspect
|
||||||
|
db.get_engine()
|
||||||
|
inspector = inspect(db.get_engine())
|
||||||
|
# monitor_id should be unique in monitors
|
||||||
|
monitors_idx = inspector.get_indexes("monitors")
|
||||||
|
assert any("monitor_id" in i.get("column_names", []) and i.get("unique") for i in monitors_idx)
|
||||||
|
# Same for actors.actor_id
|
||||||
|
actors_idx = inspector.get_indexes("actors")
|
||||||
|
assert any("actor_id" in i.get("column_names", []) and i.get("unique") for i in actors_idx)
|
||||||
Loading…
Add table
Add a link
Reference in a new issue