diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..d4a8397 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,136 @@ +name: SUQL liveness check + +# Runs the full pipeline on every push and PR: bring up Postgres, fetch the +# sample ACLED CSV from a private fixtures repo, ingest it, start SUQL's two +# servers, then issue one canned `answer()` query and assert it returns rows. +# +# Required GitHub Secrets on this fork: +# FIXTURES_TOKEN fine-grained PAT with `Contents: Read` on the private +# fixtures repo (see tests/SETUP.md for one-time setup) +# OPENAI_API_KEY LLM proxy credential +# OPENAI_API_BASE proxy base URL (e.g. https://azureopenai.genie.stanford.edu/) +# +# Optional GitHub Variable: +# FIXTURES_REPO owner/name of the private repo holding the sample CSV +# (defaults to skyxiath/suql-test-fixtures) + +on: + push: + branches: [main, add-testing-framework] + workflow_dispatch: + # No pull_request trigger: the workflow needs secrets that can't be + # available on cross-repo PRs from forks anyway. Maintainers who fork + # and want CI on their PRs should add `pull_request:` here once their + # FIXTURES_TOKEN secret is configured. + +jobs: + liveness: + runs-on: ubuntu-latest + timeout-minutes: 15 + + services: + postgres: + image: postgres:16 + env: + POSTGRES_USER: oval + POSTGRES_PASSWORD: oval + POSTGRES_DB: acled + ports: ['5432:5432'] + options: >- + --health-cmd "pg_isready -U oval" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + + env: + PGHOST: 127.0.0.1 + PGPORT: '5432' + PGDATABASE: acled + PGUSER: oval + PGPASSWORD: oval + + steps: + - name: Checkout SUQL repo + uses: actions/checkout@v4 + + - name: Fetch sample CSV from private fixtures repo + uses: actions/checkout@v4 + with: + repository: ${{ vars.FIXTURES_REPO || 'skyxiath/suql-test-fixtures' }} + token: ${{ secrets.FIXTURES_TOKEN }} + path: _fixtures + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + cache: pip + + - name: Cache HuggingFace model (BAAI/bge-large-en-v1.5, ~1.3GB) + uses: actions/cache@v4 + with: + path: ~/.cache/huggingface + key: hf-bge-large-en-v1.5 + + - name: Install SUQL with embedding deps + run: | + pip install --upgrade pip + pip install -e .[embedding] + pip install python-dotenv + + - name: Install postgresql-client + run: sudo apt-get update -qq && sudo apt-get install -y -qq postgresql-client + + - name: Ingest sample data + env: + SAMPLE_CSV: ${{ github.workspace }}/_fixtures/acled_sample.csv + run: ./tests/ingest.sh + + - name: Start free-text server (port 8500) + run: | + nohup python -m suql.free_text_fcns_server > /tmp/freetext.log 2>&1 & + echo "$!" > /tmp/freetext.pid + + - name: Start embedding server (port 8505) + # Embedding server connects to PG, reads notes, computes embeddings. + # First run on this runner downloads BAAI/bge-large-en-v1.5 (~60s); + # subsequent runs hit the HF cache. CPU mode is auto-detected. + run: | + export SUQL_EMBED_PORT=8505 + nohup python tests/start_embedding_server.py > /tmp/embed.log 2>&1 & + echo "$!" > /tmp/embed.pid + + - name: Wait for SUQL servers + run: | + for port in 8500 8505; do + for i in $(seq 1 60); do + if (echo > /dev/tcp/127.0.0.1/$port) 2>/dev/null; then + echo "✓ port $port open after ${i}s"; break + fi + sleep 2 + done + if ! (echo > /dev/tcp/127.0.0.1/$port) 2>/dev/null; then + echo "✗ port $port never came up. Logs:" + echo "--- freetext.log ---"; tail -50 /tmp/freetext.log || true + echo "--- embed.log ---"; tail -50 /tmp/embed.log || true + exit 1 + fi + done + + - name: Run liveness probe + env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + OPENAI_API_BASE: ${{ secrets.OPENAI_API_BASE }} + SUQL_EMBED_URL: http://127.0.0.1:8505 + SUQL_FREETEXT_URL: http://127.0.0.1:8500 + run: python tests/check_alive.py + + - name: Upload server logs (always, for debugging) + if: always() + uses: actions/upload-artifact@v4 + with: + name: server-logs + path: | + /tmp/freetext.log + /tmp/embed.log + retention-days: 7 diff --git a/.gitignore b/.gitignore index 9e990ad..de0f749 100644 --- a/.gitignore +++ b/.gitignore @@ -154,4 +154,10 @@ tests/private/ # Local Gradio demo — OVAL/Stanford-specific defaults, not for the public repo demo/ -src/suql/loaders/documents.py \ No newline at end of file +src/suql/loaders/documents.py + +# Sample data lives in a separate private fixtures repo (see tests/SETUP.md). +# These belt-and-suspenders entries make it impossible to accidentally commit +# the plaintext CSV into this repo even if someone drops one in for local dev. +tests/fixtures/*.csv +_fixtures/ diff --git a/setup.py b/setup.py index 6b89d86..dd86372 100644 --- a/setup.py +++ b/setup.py @@ -32,6 +32,13 @@ "FlagEmbedding~=1.2.5", ] +install_embedding_requires = [ + # ~=1.2.5 included a version with FlagEmbedding/BGE_M3/trainer.py + # referencing `Optional` without importing it — upstream-fixed in 1.3+. + "FlagEmbedding>=1.3", + "faiss-cpu>=1.7.4", +] + # Additional package information classifiers = [ "License :: OSI Approved :: Apache Software License", @@ -54,7 +61,10 @@ packages=packages, package_dir={"": "src"}, install_requires=install_requires, - extra_requires={"dev": install_dev_requires}, + extras_require={ + "dev": install_dev_requires, + "embedding": install_embedding_requires, + }, url=url, classifiers=classifiers, package_data={"": ["*.prompt"]}, diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..18fe6b4 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,114 @@ +# SUQL liveness pipeline + +End-to-end CI check: bring up Postgres, ingest a private data sample, +start SUQL's two servers, run one canned `answer()` query, confirm it +returns rows. Runs on every push via `.github/workflows/test.yml`. + +The answer to "did this change break anything obvious?" is a green or +red dot on the commit — no local environment needed to find out. + +## Layout + +``` +tests/ +├── README.md this file +├── SETUP.md one-time fixtures repo + PAT setup +├── check_alive.py liveness probe — runs the canned query +├── ingest.sh applies schema + \COPY (auth-agnostic) +├── start_embedding_server.py embedding-server bootstrap +└── fixtures/ + └── schema.sql example events table + select_user / creator_role +``` + +No CSV ever lives in this repo. The data sits in a separate private +fixtures repo that CI fetches with a scoped credential. See `SETUP.md`. + +## Required GitHub secrets + +| Secret | What it is | +| ------------------ | ---------------------------------------------------------- | +| `FIXTURES_TOKEN` | fine-grained PAT, `Contents: Read` on the fixtures repo | +| `OPENAI_API_KEY` | LLM proxy credential | +| `OPENAI_API_BASE` | proxy base URL | + +## Optional GitHub variable + +| Variable | Default | What it is | +| ----------------- | ----------------------------- | -------------------------------- | +| `FIXTURES_REPO` | `skyxiath/suql-test-fixtures` | `owner/name` of the fixtures repo | + +Set these at: **repo → Settings → Secrets and variables → Actions**. + +See `SETUP.md` for the walkthrough. + +## Adapting the probe to your data + +The example schema + canned query assume an ACLED-shaped events table. +For a different dataset: + +1. Replace `fixtures/schema.sql` with your table DDL. Keep the + `select_user` and `creator_role` blocks — SUQL's compiler needs them. +2. Push a CSV matching that schema to the fixtures repo. +3. Update `start_embedding_server.py`'s call to `store.add(...)` to + point at your table / primary key / free-text column. Or override + via `SUQL_TABLE`, `SUQL_ID_COL`, `SUQL_TEXT_COL` env vars. +4. Override the probe via `SUQL_QUERY` in the workflow's "Run liveness + probe" step. Any query that exercises one `answer()` clause and + returns rows from your sample works. + +## Refreshing the sample + +Sample lives in the fixtures repo, not here. Update it there: + +```bash +cd /path/to/your-fixtures-repo +# (generate / copy in a new sample.csv however you do) +git add sample.csv && git commit -m "Refresh sample" && git push +``` + +CI on the next workflow run picks up the new sample automatically — no +change to this repo needed. + +## Running locally + +```bash +# 1. Start Postgres locally: +docker run -d --name suql-test-pg -p 5432:5432 \ + -e POSTGRES_USER=oval -e POSTGRES_PASSWORD=oval -e POSTGRES_DB=acled \ + postgres:16 + +# 2. Install SUQL with embedding deps: +pip install -e .[embedding] +pip install python-dotenv + +# 3. Drop a sample CSV at /tmp/sample.csv (from wherever you keep it). + +# 4. Ingest: +PGPASSWORD=oval SAMPLE_CSV=/tmp/sample.csv ./tests/ingest.sh + +# 5. Start servers: +python -m suql.free_text_fcns_server & +SUQL_EMBED_PORT=8505 python tests/start_embedding_server.py & + +# 6. Run probe: +OPENAI_API_KEY=... OPENAI_API_BASE=... python tests/check_alive.py +``` + +## What "liveness" means here + +The probe runs one canned `answer(...) = 'yes'` query. Exit codes: + +| Code | Meaning | +| ---- | ----------------------------------------------------------- | +| 0 | `suql_execute` returned >= 1 row | +| 1 | `suql_execute` raised (PG, embedding, free-text, or LLM) | +| 2 | Ran cleanly but returned 0 rows | +| 3 | Preflight failed — TCP target unreachable or creds missing | + +The workflow fails on any non-zero exit. + +## Scope + +This catches *plumbing* breaks (server boot failure, SQL compilation, +prompt loading, network connectivity). It does not measure result +quality — for that you want benchmarks / evals, not a CI smoke test. diff --git a/tests/SETUP.md b/tests/SETUP.md new file mode 100644 index 0000000..379e214 --- /dev/null +++ b/tests/SETUP.md @@ -0,0 +1,105 @@ +# One-time setup: fixtures repo + PAT + secrets + +This walkthrough wires up the auth-gated fetch the CI workflow depends on. +Do it once; CI handles itself from there. + +Notation: `` is your GitHub username/org. `` is the +name you'll pick for the data repo (e.g. `suql-test-fixtures`). + +## Step 1 — Create the private fixtures repo + +The sample CSV needs a home separate from the SUQL code. Create a new +**private** repo: + +```bash +gh repo create / --private --description \ + "Sample data + fixtures consumed by SUQL CI" +``` + +That's it for now — it can stay empty until you push the first sample. + +## Step 2 — Drop in the first sample + +Push a CSV named (by default) `acled_sample.csv` matching the column +order in `tests/fixtures/schema.sql`: + +```bash +gh repo clone / /tmp/fixtures +cp /path/to/your/sample.csv /tmp/fixtures/acled_sample.csv +cd /tmp/fixtures +git add acled_sample.csv +git commit -m "Initial sample" +git push +``` + +For a different filename, set the `SAMPLE_CSV` env in the workflow's +"Ingest sample data" step to match. + +## Step 3 — Create a fine-grained PAT + +GitHub's fine-grained PATs let you scope a token to a single repo with +read-only permission. That's exactly what CI needs. + +1. Open: **GitHub → Settings → Developer settings → Personal access tokens + → Fine-grained tokens → Generate new token** +2. **Token name:** something memorable, e.g. `suql-ci-fixtures-readonly` +3. **Expiration:** pick the longest you're comfortable with (max 1 year). + Set a calendar reminder to rotate before expiry. +4. **Repository access:** *Only select repositories* → `/` +5. **Permissions → Repository permissions:** + - `Contents`: **Read-only** + - `Metadata`: **Read-only** (required by GitHub for any access) +6. Click **Generate token**, copy the value — GitHub only shows it once. + +## Step 4 — Add secrets + variable to the SUQL repo + +1. Open: **SUQL fork → Settings → Secrets and variables → Actions → + Secrets tab → New repository secret** +2. Add three secrets: + +| Name | Value | +| ----------------- | ---------------------------------------------- | +| `FIXTURES_TOKEN` | the PAT from Step 3 | +| `OPENAI_API_KEY` | your LLM proxy key | +| `OPENAI_API_BASE` | your LLM proxy base URL | + +3. Switch to the **Variables** tab and add one: + +| Name | Value | +| ---------------- | ---------------------------------- | +| `FIXTURES_REPO` | `/` | + +## Step 5 — Verify by triggering a CI run + +```bash +git push +# or +gh workflow run "SUQL liveness check" +``` + +Watch: **repo → Actions → SUQL liveness check → most-recent run**. + +The "Fetch sample CSV" step is where the PAT is exercised. If it fails: +- 404 → PAT can see fewer repos than expected; recheck Step 3's + *Repository access* selection. +- 403 → PAT lacks `Contents: Read`; recheck Step 3's *Permissions*. +- Anything else → check the run logs. + +The "Run liveness probe" step is the actual end-to-end check. On +success it prints `✓ suql_execute returned in Ns` with row count and +cost. + +## Rotation + +When the PAT approaches expiration (GitHub emails you ahead of time): + +1. Generate a new PAT with the same scope (Step 3) +2. Update the `FIXTURES_TOKEN` secret with the new value (Step 4) +3. Delete the old PAT + +## Revocation (if the PAT leaks) + +1. **Immediately** delete the leaked PAT at GitHub → Settings → Developer + settings → Personal access tokens → Fine-grained tokens +2. Generate a new PAT, update the secret, push a no-op commit to verify +3. Review whatever process exposed the PAT diff --git a/tests/check_alive.py b/tests/check_alive.py new file mode 100755 index 0000000..83fce0c --- /dev/null +++ b/tests/check_alive.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python +"""SUQL liveness probe — one query, one verdict. + +The whole point of this script: prove that a SUQL stack — Postgres, +embedding server, free-text server, LLM proxy — is end-to-end functional. +It is the script GitHub Actions runs at the end of the pipeline to decide +whether the build is green. + +Every connection target is an env var, so the same script runs unchanged +locally and in CI. + +Exit codes: + 0 SUQL returned >= 1 row + 1 SUQL raised — connection, server, LLM, or compiler error + 2 SUQL ran cleanly but returned 0 rows (alive but suspicious) + 3 preflight failed — a TCP target was unreachable or creds missing + +Usage:: + + python tests/check_alive.py # use env / defaults + python tests/check_alive.py --print-env # show resolved config, exit + python tests/check_alive.py --no-preflight # skip TCP probe, fail inside SUQL instead +""" + +from __future__ import annotations + +import argparse +import os +import socket +import sys +import time +from pathlib import Path +from urllib.parse import urlparse + +# --------------------------------------------------------------------------- +# Bootstrap: load .env if present, ensure suql imports from this repo's src/ +# --------------------------------------------------------------------------- +_HERE = Path(__file__).resolve().parent +_REPO_ROOT = _HERE.parent + +try: + from dotenv import load_dotenv + + env_path = _REPO_ROOT / ".env" + if env_path.exists(): + load_dotenv(env_path) +except ImportError: + pass # In CI, env vars come from the workflow directly + +_SRC = _REPO_ROOT / "src" +if _SRC.exists(): + sys.path.insert(0, str(_SRC)) + + +# Defaults assume the workflow's service containers + the example fixtures +# schema. Override any of these from the calling environment for a different +# database, table layout, or LLM proxy. +DEFAULTS = { + "PGHOST": "127.0.0.1", + "PGPORT": "5432", + "PGDATABASE": "acled", + "PGUSER": "oval", + "PGPASSWORD": "oval", + "SUQL_EMBED_URL": "http://127.0.0.1:8505", + "SUQL_FREETEXT_URL": "http://127.0.0.1:8500", + "SUQL_SELECT_USER": "select_user", + "SUQL_SELECT_PSWD": "select_user", + "SUQL_CREATE_USER": "creator_role", + "SUQL_CREATE_PSWD": "creator_role", + "SUQL_TABLE": "events", + "SUQL_ID_COL": "event_id_cnty", + "SUQL_TIMEOUT_MS": "120000", +} + +# Override with SUQL_QUERY env var. The default below pairs with the example +# schema in tests/fixtures/schema.sql; any query that exercises one answer() +# clause and returns >= 1 row from your data works. +DEFAULT_QUERY = """\ +SELECT event_id_cnty, country, event_type, notes +FROM events +WHERE notes IS NOT NULL + AND answer(notes, 'Does this describe a conflict, protest, or violent event?') = 'yes' +LIMIT 3; +""" + + +def _env(key: str) -> str: + return os.environ.get(key, DEFAULTS.get(key, "")) + + +def _port_open(host: str, port: int, timeout: float = 1.0) -> bool: + try: + with socket.create_connection((host, port), timeout=timeout): + return True + except OSError: + return False + + +def _preflight() -> list[str]: + problems: list[str] = [] + if not _port_open(_env("PGHOST"), int(_env("PGPORT"))): + problems.append(f"Postgres unreachable at {_env('PGHOST')}:{_env('PGPORT')}") + for label, key in [("embedding", "SUQL_EMBED_URL"), ("free-text", "SUQL_FREETEXT_URL")]: + u = urlparse(_env(key)) + port = u.port or (443 if u.scheme == "https" else 80) + if not _port_open(u.hostname or "127.0.0.1", port): + problems.append(f"SUQL {label} server unreachable at {_env(key)}") + if not os.environ.get("OPENAI_API_KEY") or not os.environ.get("OPENAI_API_BASE"): + problems.append("OPENAI_API_KEY / OPENAI_API_BASE not set") + return problems + + +def main() -> int: + parser = argparse.ArgumentParser(description="SUQL liveness probe") + parser.add_argument("--print-env", action="store_true", + help="Show resolved config and exit.") + parser.add_argument("--no-preflight", action="store_true", + help="Skip TCP / cred preflight; let SUQL fail instead.") + args = parser.parse_args() + + sql = os.environ.get("SUQL_QUERY", DEFAULT_QUERY).strip() + config = {k: _env(k) for k in DEFAULTS} + config["OPENAI_API_BASE"] = os.environ.get("OPENAI_API_BASE", "") + config["OPENAI_API_KEY"] = "" if os.environ.get("OPENAI_API_KEY") else "" + + print("=== SUQL liveness probe ===") + for k, v in config.items(): + print(f" {k:18s} {v}") + print() + print("Query:") + for line in sql.splitlines(): + print(f" {line}") + print() + + if args.print_env: + return 0 + + if not args.no_preflight: + problems = _preflight() + if problems: + print("✗ Preflight failed:") + for p in problems: + print(f" - {p}") + return 3 + print("✓ Preflight: all ports open, API creds set") + print() + + try: + from suql import suql_execute + except ImportError as e: + print(f"✗ Cannot import suql: {e}") + return 1 + + t0 = time.time() + try: + rows, cols, cache = suql_execute( + sql, + table_w_ids={_env("SUQL_TABLE"): _env("SUQL_ID_COL")}, + database=_env("PGDATABASE"), + embedding_server_address=_env("SUQL_EMBED_URL"), + free_text_server_address=_env("SUQL_FREETEXT_URL"), + host=_env("PGHOST"), + port=int(_env("PGPORT")), + select_username=_env("SUQL_SELECT_USER"), + select_userpswd=_env("SUQL_SELECT_PSWD"), + create_username=_env("SUQL_CREATE_USER"), + create_userpswd=_env("SUQL_CREATE_PSWD"), + statement_timeout=int(_env("SUQL_TIMEOUT_MS")), + api_base=os.environ.get("OPENAI_API_BASE"), + api_key=os.environ.get("OPENAI_API_KEY"), + ) + except Exception as e: + elapsed = time.time() - t0 + print(f"✗ suql_execute raised after {elapsed:.1f}s") + print(f" {type(e).__name__}: {e}") + return 1 + + elapsed = time.time() - t0 + stats = (cache or {}).get("_stats", {}) + print(f"✓ suql_execute returned in {elapsed:.1f}s") + print(f" rows: {len(rows)}") + print(f" cols: {cols}") + print(f" cost: ${stats.get('cost', 0):.4f}") + print(f" calls: {stats.get('calls', 0)}") + if rows: + first = rows[0] + preview = [(str(v)[:80] + "…") if len(str(v)) > 80 else str(v) for v in first] + print(f" row[0]: {preview}") + return 0 + + print() + print("⚠ Query ran cleanly but returned zero rows.") + print(" SUQL is alive but either (a) the LLM judged every candidate as 'no',") + print(" (b) the WHERE clause filtered everything out, or (c) the embedding") + print(" server returned no neighbours. Inspect cache for details.") + return 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/fixtures/schema.sql b/tests/fixtures/schema.sql new file mode 100644 index 0000000..ab6e821 --- /dev/null +++ b/tests/fixtures/schema.sql @@ -0,0 +1,45 @@ +-- Example schema for the CI liveness pipeline. Replace with your own +-- DDL for any dataset that has at least one free-text column SUQL can +-- run answer() over. +-- +-- Two things to preserve when swapping schemas: +-- 1. The select_user / creator_role pair below. SUQL's compiler needs +-- a distinct creator role for its temp-table dance. +-- 2. The column-list ordering of the CSV pushed to the fixtures repo +-- must match the CREATE TABLE column order, since the workflow +-- uses \COPY ... WITH HEADER. + +CREATE TABLE IF NOT EXISTS events ( + event_id_cnty TEXT PRIMARY KEY, + event_date DATE NOT NULL, + year INTEGER, + event_type TEXT, + sub_event_type TEXT, + country TEXT, + admin1 TEXT, + admin2 TEXT, + admin3 TEXT, + location TEXT, + fatalities INTEGER, + notes TEXT +); + +-- SUQL's compiler needs roles distinct from the connecting user for its +-- temp-table dance (creator_role creates tables, select_user reads). We +-- create both here so the workflow can SET ROLE into them without GRANT +-- ceremony per query. +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'select_user') THEN + CREATE ROLE select_user LOGIN PASSWORD 'select_user'; + END IF; + IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'creator_role') THEN + CREATE ROLE creator_role LOGIN PASSWORD 'creator_role'; + END IF; +END$$; + +GRANT SELECT ON events TO select_user; +GRANT ALL ON SCHEMA public TO creator_role; +GRANT CREATE ON SCHEMA public TO creator_role; + +CREATE INDEX IF NOT EXISTS events_country_date_idx ON events (country, event_date); diff --git a/tests/ingest.sh b/tests/ingest.sh new file mode 100755 index 0000000..c1ba025 --- /dev/null +++ b/tests/ingest.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# Apply the schema and \COPY a sample CSV into Postgres. +# +# This script is auth-agnostic: it expects the CSV to already exist on disk at +# $SAMPLE_CSV. Fetching the CSV (from a private fixtures repo, an S3 bucket, +# wherever) is the *caller's* job — keeps secrets handling out of this script. +# +# In CI, the workflow checks out the fixtures repo first and points SAMPLE_CSV +# at the resulting file. Locally, you drop a CSV anywhere and export the path. +# +# Env vars: +# SAMPLE_CSV (required) absolute path to the CSV to load +# PGHOST default 127.0.0.1 +# PGPORT default 5432 +# PGDATABASE default acled +# PGUSER default oval +# PGPASSWORD (required, passed via env not flag) + +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SCHEMA="$HERE/fixtures/schema.sql" + +: "${SAMPLE_CSV:?SAMPLE_CSV must point at a CSV file on disk}" +: "${PGHOST:=127.0.0.1}" +: "${PGPORT:=5432}" +: "${PGDATABASE:=acled}" +: "${PGUSER:=oval}" +: "${PGPASSWORD:?PGPASSWORD must be set}" + +export PGPASSWORD + +[ -r "$SCHEMA" ] || { echo "✗ schema not found: $SCHEMA" >&2; exit 1; } +[ -r "$SAMPLE_CSV" ] || { echo "✗ sample CSV not found: $SAMPLE_CSV" >&2; exit 1; } + +ROW_COUNT="$(wc -l < "$SAMPLE_CSV" | tr -d ' ')" +echo "→ Sample CSV: $SAMPLE_CSV ($ROW_COUNT lines incl. header)" + +echo "→ Applying schema..." +psql -v ON_ERROR_STOP=1 -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -d "$PGDATABASE" \ + -f "$SCHEMA" > /dev/null + +echo "→ Loading CSV..." +psql -v ON_ERROR_STOP=1 -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -d "$PGDATABASE" \ + -c "\\COPY events FROM '$SAMPLE_CSV' WITH (FORMAT csv, HEADER true)" + +INGESTED="$(psql -t -A -h "$PGHOST" -p "$PGPORT" -U "$PGUSER" -d "$PGDATABASE" \ + -c "SELECT COUNT(*) FROM events;")" +echo "✓ events table contains $INGESTED rows" diff --git a/tests/start_embedding_server.py b/tests/start_embedding_server.py new file mode 100755 index 0000000..373ead3 --- /dev/null +++ b/tests/start_embedding_server.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python +"""Start the SUQL embedding server pointed at the ACLED ``events`` table. + +SUQL's bundled ``python -m suql.faiss_embedding`` is hard-coded for the +restaurants demo (see ``faiss_embedding.py``'s ``__main__``). This script +is the ACLED equivalent: register ``events.notes`` with a +``MultipleEmbeddingStore`` and start the HTTP server. + +Reads from env (with sensible defaults that match the CI workflow): + PGDATABASE (default: acled) + PGUSER (default: oval) + PGPASSWORD (default: oval) + SUQL_EMBED_PORT (default: 8505) + SUQL_TABLE (default: events) + SUQL_ID_COL (default: event_id_cnty) + SUQL_TEXT_COL (default: notes) +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +# Make sure we import the checked-out suql, not a stray pip-installed copy. +_REPO_ROOT = Path(__file__).resolve().parent.parent +_SRC = _REPO_ROOT / "src" +if _SRC.exists(): + sys.path.insert(0, str(_SRC)) + +from suql.faiss_embedding import MultipleEmbeddingStore + + +def main() -> None: + store = MultipleEmbeddingStore() + store.add( + table_name=os.environ.get("SUQL_TABLE", "events"), + primary_key_field_name=os.environ.get("SUQL_ID_COL", "event_id_cnty"), + free_text_field_name=os.environ.get("SUQL_TEXT_COL", "notes"), + db_name=os.environ.get("PGDATABASE", "acled"), + user=os.environ.get("PGUSER", "oval"), + password=os.environ.get("PGPASSWORD", "oval"), + ) + store.start_embedding_server( + host="127.0.0.1", + port=int(os.environ.get("SUQL_EMBED_PORT", "8505")), + ) + + +if __name__ == "__main__": + main()