Upload notebooks/00_setup/02_schema_migration.ipynb with huggingface_hub

notebooks/00_setup/02_schema_migration.ipynb

@@ -0,0 +1,298 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Schema Migration — Analysis Tables\n",
+    "\n",
+    "Creates all database tables required by the analysis pipelines.\n",
+    "This notebook is idempotent (`CREATE TABLE IF NOT EXISTS`) and safe to re-run.\n",
+    "\n",
+    "Supports Papermill parameterization: set `dry_run = True` to print DDL without executing."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {
+    "tags": [
+     "parameters"
+    ]
+   },
+   "source": [
+    "# Parameters\n",
+    "dry_run = False"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "from research_lib.db import fetch_df\n",
+    "from research_lib.config import get_engine\n",
+    "\n",
+    "engine = get_engine()\n",
+    "\n",
+    "DDL_STATEMENTS = [\n",
+    "    # --- analysis_runs ---\n",
+    "    \"\"\"\n",
+    "    CREATE TABLE IF NOT EXISTS analysis_runs (\n",
+    "        id SERIAL PRIMARY KEY,\n",
+    "        pipeline_name TEXT NOT NULL,\n",
+    "        source_section TEXT,\n",
+    "        started_at TIMESTAMP NOT NULL DEFAULT NOW(),\n",
+    "        finished_at TIMESTAMP,\n",
+    "        status TEXT NOT NULL DEFAULT 'running',\n",
+    "        documents_processed INTEGER DEFAULT 0,\n",
+    "        parameters JSONB,\n",
+    "        error_message TEXT\n",
+    "    );\n",
+    "    \"\"\",\n",
+    "    # --- entity_aliases ---\n",
+    "    \"\"\"\n",
+    "    CREATE TABLE IF NOT EXISTS entity_aliases (\n",
+    "        id SERIAL PRIMARY KEY,\n",
+    "        canonical_name TEXT NOT NULL,\n",
+    "        alias_name TEXT NOT NULL,\n",
+    "        entity_type TEXT NOT NULL,\n",
+    "        confidence REAL DEFAULT 1.0,\n",
+    "        created_at TIMESTAMP DEFAULT NOW(),\n",
+    "        UNIQUE(alias_name, entity_type)\n",
+    "    );\n",
+    "    \"\"\",\n",
+    "    # --- entity_relationships ---\n",
+    "    \"\"\"\n",
+    "    CREATE TABLE IF NOT EXISTS entity_relationships (\n",
+    "        id SERIAL PRIMARY KEY,\n",
+    "        entity_a TEXT NOT NULL,\n",
+    "        entity_a_type TEXT NOT NULL,\n",
+    "        entity_b TEXT NOT NULL,\n",
+    "        entity_b_type TEXT NOT NULL,\n",
+    "        co_occurrence_count INTEGER NOT NULL DEFAULT 1,\n",
+    "        document_count INTEGER NOT NULL DEFAULT 1,\n",
+    "        source_section TEXT,\n",
+    "        avg_distance REAL,\n",
+    "        sample_doc_ids INTEGER[],\n",
+    "        created_at TIMESTAMP DEFAULT NOW(),\n",
+    "        UNIQUE(entity_a, entity_a_type, entity_b, entity_b_type, source_section)\n",
+    "    );\n",
+    "    \"\"\",\n",
+    "    # --- topics ---\n",
+    "    \"\"\"\n",
+    "    CREATE TABLE IF NOT EXISTS topics (\n",
+    "        id SERIAL PRIMARY KEY,\n",
+    "        model_name TEXT NOT NULL,\n",
+    "        source_section TEXT,\n",
+    "        topic_id INTEGER NOT NULL,\n",
+    "        topic_label TEXT,\n",
+    "        top_words TEXT[] NOT NULL,\n",
+    "        top_word_scores REAL[],\n",
+    "        document_count INTEGER,\n",
+    "        representative_doc_ids INTEGER[],\n",
+    "        embedding vector(384),\n",
+    "        created_at TIMESTAMP DEFAULT NOW(),\n",
+    "        UNIQUE(model_name, source_section, topic_id)\n",
+    "    );\n",
+    "    \"\"\",\n",
+    "    # --- document_topics ---\n",
+    "    \"\"\"\n",
+    "    CREATE TABLE IF NOT EXISTS document_topics (\n",
+    "        id SERIAL PRIMARY KEY,\n",
+    "        document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,\n",
+    "        model_name TEXT NOT NULL,\n",
+    "        topic_id INTEGER NOT NULL,\n",
+    "        probability REAL,\n",
+    "        created_at TIMESTAMP DEFAULT NOW(),\n",
+    "        UNIQUE(document_id, model_name)\n",
+    "    );\n",
+    "    \"\"\",\n",
+    "    # --- document_features ---\n",
+    "    \"\"\"\n",
+    "    CREATE TABLE IF NOT EXISTS document_features (\n",
+    "        id SERIAL PRIMARY KEY,\n",
+    "        document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,\n",
+    "        feature_name TEXT NOT NULL,\n",
+    "        feature_value REAL,\n",
+    "        feature_json JSONB,\n",
+    "        created_at TIMESTAMP DEFAULT NOW(),\n",
+    "        UNIQUE(document_id, feature_name)\n",
+    "    );\n",
+    "    \"\"\",\n",
+    "    # --- page_features ---\n",
+    "    \"\"\"\n",
+    "    CREATE TABLE IF NOT EXISTS page_features (\n",
+    "        id SERIAL PRIMARY KEY,\n",
+    "        page_id INTEGER NOT NULL REFERENCES pages(id) ON DELETE CASCADE,\n",
+    "        feature_name TEXT NOT NULL,\n",
+    "        feature_value REAL,\n",
+    "        feature_json JSONB,\n",
+    "        created_at TIMESTAMP DEFAULT NOW(),\n",
+    "        UNIQUE(page_id, feature_name)\n",
+    "    );\n",
+    "    \"\"\",\n",
+    "    # --- duplicate_pairs ---\n",
+    "    \"\"\"\n",
+    "    CREATE TABLE IF NOT EXISTS duplicate_pairs (\n",
+    "        id SERIAL PRIMARY KEY,\n",
+    "        page_id_a INTEGER NOT NULL REFERENCES pages(id) ON DELETE CASCADE,\n",
+    "        page_id_b INTEGER NOT NULL REFERENCES pages(id) ON DELETE CASCADE,\n",
+    "        similarity REAL NOT NULL,\n",
+    "        method TEXT NOT NULL DEFAULT 'embedding',\n",
+    "        created_at TIMESTAMP DEFAULT NOW(),\n",
+    "        UNIQUE(page_id_a, page_id_b, method)\n",
+    "    );\n",
+    "    \"\"\",\n",
+    "    # --- document_keywords ---\n",
+    "    \"\"\"\n",
+    "    CREATE TABLE IF NOT EXISTS document_keywords (\n",
+    "        id SERIAL PRIMARY KEY,\n",
+    "        document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,\n",
+    "        keyword TEXT NOT NULL,\n",
+    "        score REAL NOT NULL,\n",
+    "        method TEXT NOT NULL DEFAULT 'tfidf',\n",
+    "        created_at TIMESTAMP DEFAULT NOW()\n",
+    "    );\n",
+    "    \"\"\",\n",
+    "]\n",
+    "\n",
+    "if dry_run:\n",
+    "    print(\"=== DRY RUN — DDL statements that would be executed ===\")\n",
+    "    for i, stmt in enumerate(DDL_STATEMENTS, 1):\n",
+    "        print(f\"\\n--- Statement {i} ---\")\n",
+    "        print(stmt.strip())\n",
+    "else:\n",
+    "    with engine.begin() as conn:\n",
+    "        for i, stmt in enumerate(DDL_STATEMENTS, 1):\n",
+    "            print(f\"Executing DDL statement {i}/{len(DDL_STATEMENTS)}...\")\n",
+    "            conn.execute(stmt)\n",
+    "    print(f\"\\nAll {len(DDL_STATEMENTS)} DDL statements executed successfully.\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Create indexes for query performance\n",
+    "\n",
+    "INDEX_STATEMENTS = [\n",
+    "    # analysis_runs\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_analysis_runs_pipeline ON analysis_runs (pipeline_name);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_analysis_runs_status ON analysis_runs (status);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_analysis_runs_section ON analysis_runs (source_section);\",\n",
+    "    # entity_aliases\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_entity_aliases_canonical ON entity_aliases (canonical_name);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_entity_aliases_type ON entity_aliases (entity_type);\",\n",
+    "    # entity_relationships\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_entity_rel_a ON entity_relationships (entity_a, entity_a_type);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_entity_rel_b ON entity_relationships (entity_b, entity_b_type);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_entity_rel_section ON entity_relationships (source_section);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_entity_rel_cooccurrence ON entity_relationships (co_occurrence_count DESC);\",\n",
+    "    # topics\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_topics_model ON topics (model_name);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_topics_section ON topics (source_section);\",\n",
+    "    # document_topics\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_doc_topics_doc ON document_topics (document_id);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_doc_topics_model ON document_topics (model_name);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_doc_topics_topic ON document_topics (topic_id);\",\n",
+    "    # document_features\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_doc_features_doc ON document_features (document_id);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_doc_features_name ON document_features (feature_name);\",\n",
+    "    # page_features\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_page_features_page ON page_features (page_id);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_page_features_name ON page_features (feature_name);\",\n",
+    "    # duplicate_pairs\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_dup_pairs_a ON duplicate_pairs (page_id_a);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_dup_pairs_b ON duplicate_pairs (page_id_b);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_dup_pairs_sim ON duplicate_pairs (similarity DESC);\",\n",
+    "    # document_keywords\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_doc_kw_doc ON document_keywords (document_id);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_doc_kw_keyword ON document_keywords (keyword);\",\n",
+    "    \"CREATE INDEX IF NOT EXISTS idx_doc_kw_method ON document_keywords (method);\",\n",
+    "]\n",
+    "\n",
+    "if dry_run:\n",
+    "    print(\"=== DRY RUN — Index statements ===\")\n",
+    "    for stmt in INDEX_STATEMENTS:\n",
+    "        print(f\"  {stmt}\")\n",
+    "else:\n",
+    "    with engine.begin() as conn:\n",
+    "        for i, stmt in enumerate(INDEX_STATEMENTS, 1):\n",
+    "            conn.execute(stmt)\n",
+    "    print(f\"Created {len(INDEX_STATEMENTS)} indexes successfully.\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Grant SELECT on all new tables to web_reader\n",
+    "\n",
+    "NEW_TABLES = [\n",
+    "    \"analysis_runs\",\n",
+    "    \"entity_aliases\",\n",
+    "    \"entity_relationships\",\n",
+    "    \"topics\",\n",
+    "    \"document_topics\",\n",
+    "    \"document_features\",\n",
+    "    \"page_features\",\n",
+    "    \"duplicate_pairs\",\n",
+    "    \"document_keywords\",\n",
+    "]\n",
+    "\n",
+    "if dry_run:\n",
+    "    print(\"=== DRY RUN — GRANT statements ===\")\n",
+    "    for table in NEW_TABLES:\n",
+    "        print(f\"  GRANT SELECT ON {table} TO web_reader;\")\n",
+    "else:\n",
+    "    with engine.begin() as conn:\n",
+    "        for table in NEW_TABLES:\n",
+    "            conn.execute(f\"GRANT SELECT ON {table} TO web_reader;\")\n",
+    "            print(f\"  GRANT SELECT ON {table} TO web_reader -- OK\")\n",
+    "    print(f\"\\nGranted SELECT to web_reader on {len(NEW_TABLES)} tables.\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Confirm: list all tables in the database\n",
+    "from research_lib.db import fetch_df\n",
+    "\n",
+    "df_tables = fetch_df(\"\"\"\n",
+    "    SELECT table_name, \n",
+    "           pg_size_pretty(pg_total_relation_size(quote_ident(table_name))) AS total_size\n",
+    "    FROM information_schema.tables\n",
+    "    WHERE table_schema = 'public'\n",
+    "    ORDER BY table_name\n",
+    "\"\"\")\n",
+    "\n",
+    "print(f\"Total tables in public schema: {len(df_tables)}\")\n",
+    "print()\n",
+    "print(df_tables.to_string(index=False))"
+   ],
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}