forked from 0xWheatyz/SPARC
Add multi-tenant support with owner_id isolation
- Add owner_id (FK to users) column to llm_messages, jobs, and
tracked_companies tables via schema migration in initialize_schema()
- Filter all read/write operations by authenticated user's owner_id
so users cannot see or modify each other's data
- Add user-scoped /tracked endpoints alongside existing admin ones
- Add admin-scoped /admin/analyses and /admin/jobs endpoints that
return cross-tenant data without owner filtering
- Create migration script (scripts/migrate_add_owner_id.py) that
backfills owner_id=1 for all existing rows
- Replace global UNIQUE on tracked_companies.company_name with
per-owner unique index (company_name, owner_id)
- Fix route ordering: /analyze/batch and /analyze/patent routes now
registered before /analyze/{company_name} to prevent path conflicts
- Update all existing API tests with proper auth headers and owner_id
assertions
- Add comprehensive cross-tenant isolation test suite
(tests/test_multi_tenant.py)
Closes leeworks-agents/SPARC#1677
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,132 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Migration: add owner_id columns and backfill existing rows.
|
||||
|
||||
This script adds an ``owner_id`` column (FK to ``users``) to the
|
||||
``llm_messages``, ``jobs``, and ``tracked_companies`` tables, then
|
||||
backfills all existing rows with ``owner_id = 1`` (the default admin user).
|
||||
|
||||
It also replaces the old global UNIQUE constraint on
|
||||
``tracked_companies.company_name`` with a per-owner unique index so that
|
||||
different users can independently track the same company.
|
||||
|
||||
Usage:
|
||||
python scripts/migrate_add_owner_id.py
|
||||
|
||||
The script is idempotent — running it multiple times is safe.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import psycopg2
|
||||
|
||||
DATABASE_URL = os.getenv(
|
||||
"DATABASE_URL",
|
||||
"postgresql://postgres:postgres@localhost:5432/sparc",
|
||||
)
|
||||
|
||||
DEFAULT_OWNER_ID = 1
|
||||
|
||||
|
||||
def run_migration():
|
||||
"""Execute the migration."""
|
||||
conn = psycopg2.connect(DATABASE_URL)
|
||||
conn.autocommit = False
|
||||
|
||||
try:
|
||||
with conn.cursor() as cur:
|
||||
# ---------- 1. Add owner_id columns if missing ----------
|
||||
cur.execute("""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM information_schema.columns
|
||||
WHERE table_name = 'llm_messages' AND column_name = 'owner_id'
|
||||
) THEN
|
||||
ALTER TABLE llm_messages ADD COLUMN owner_id INTEGER REFERENCES users(id);
|
||||
END IF;
|
||||
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM information_schema.columns
|
||||
WHERE table_name = 'jobs' AND column_name = 'owner_id'
|
||||
) THEN
|
||||
ALTER TABLE jobs ADD COLUMN owner_id INTEGER REFERENCES users(id);
|
||||
END IF;
|
||||
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM information_schema.columns
|
||||
WHERE table_name = 'tracked_companies' AND column_name = 'owner_id'
|
||||
) THEN
|
||||
ALTER TABLE tracked_companies ADD COLUMN owner_id INTEGER REFERENCES users(id);
|
||||
END IF;
|
||||
END $$;
|
||||
""")
|
||||
|
||||
# ---------- 2. Backfill owner_id = DEFAULT_OWNER_ID ----------
|
||||
cur.execute(
|
||||
"UPDATE llm_messages SET owner_id = %s WHERE owner_id IS NULL",
|
||||
(DEFAULT_OWNER_ID,),
|
||||
)
|
||||
messages_updated = cur.rowcount
|
||||
print(f" llm_messages: backfilled {messages_updated} rows")
|
||||
|
||||
cur.execute(
|
||||
"UPDATE jobs SET owner_id = %s WHERE owner_id IS NULL",
|
||||
(DEFAULT_OWNER_ID,),
|
||||
)
|
||||
jobs_updated = cur.rowcount
|
||||
print(f" jobs: backfilled {jobs_updated} rows")
|
||||
|
||||
cur.execute(
|
||||
"UPDATE tracked_companies SET owner_id = %s WHERE owner_id IS NULL",
|
||||
(DEFAULT_OWNER_ID,),
|
||||
)
|
||||
tracked_updated = cur.rowcount
|
||||
print(f" tracked_companies: backfilled {tracked_updated} rows")
|
||||
|
||||
# ---------- 3. Create indexes ----------
|
||||
cur.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_owner
|
||||
ON llm_messages(owner_id)
|
||||
""")
|
||||
cur.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_jobs_owner
|
||||
ON jobs(owner_id)
|
||||
""")
|
||||
cur.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_tracked_companies_owner
|
||||
ON tracked_companies(owner_id)
|
||||
""")
|
||||
|
||||
# ---------- 4. Replace unique constraint on tracked_companies ----------
|
||||
cur.execute("""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (
|
||||
SELECT 1 FROM pg_constraint
|
||||
WHERE conname = 'tracked_companies_company_name_key'
|
||||
) THEN
|
||||
ALTER TABLE tracked_companies
|
||||
DROP CONSTRAINT tracked_companies_company_name_key;
|
||||
END IF;
|
||||
END $$;
|
||||
""")
|
||||
cur.execute("""
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_tracked_company_owner
|
||||
ON tracked_companies(LOWER(company_name), owner_id)
|
||||
""")
|
||||
|
||||
conn.commit()
|
||||
print("Migration completed successfully.")
|
||||
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
print("Migration FAILED — rolled back.", file=sys.stderr)
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(f"Running owner_id migration against {DATABASE_URL.split('@')[-1]} ...")
|
||||
run_migration()
|
||||
Reference in New Issue
Block a user