Initial scaffold: .gitignore, config, requirements, data seeds, directory structure
This commit is contained in:
n0mad1k
@@ -0,0 +1,191 @@
|
||||
"""Schema versioning and migration runner for Mosaic database."""
|
||||
import sqlite3
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MIGRATIONS = [
|
||||
# Version 1: Initial schema
|
||||
{
|
||||
"version": 1,
|
||||
"description": "Initial schema",
|
||||
"sql": [
|
||||
"""CREATE TABLE IF NOT EXISTS documents (
|
||||
id TEXT PRIMARY KEY,
|
||||
source TEXT NOT NULL DEFAULT '',
|
||||
source_url TEXT NOT NULL DEFAULT '',
|
||||
doc_type TEXT NOT NULL DEFAULT '',
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
date TEXT,
|
||||
classification TEXT,
|
||||
raw_path TEXT NOT NULL DEFAULT '',
|
||||
text TEXT NOT NULL DEFAULT '',
|
||||
metadata TEXT NOT NULL DEFAULT '{}',
|
||||
content_hash TEXT NOT NULL DEFAULT '',
|
||||
fetch_date TEXT NOT NULL DEFAULT '',
|
||||
etag TEXT,
|
||||
last_modified TEXT,
|
||||
char_count INTEGER NOT NULL DEFAULT 0,
|
||||
status TEXT NOT NULL DEFAULT 'PENDING',
|
||||
parse_status TEXT NOT NULL DEFAULT 'PENDING',
|
||||
last_extracted_at TEXT
|
||||
)""",
|
||||
"""CREATE INDEX IF NOT EXISTS idx_documents_source ON documents(source)""",
|
||||
"""CREATE INDEX IF NOT EXISTS idx_documents_status ON documents(status)""",
|
||||
"""CREATE INDEX IF NOT EXISTS idx_documents_parse_status ON documents(parse_status)""",
|
||||
"""CREATE INDEX IF NOT EXISTS idx_documents_date ON documents(date)""",
|
||||
"""CREATE INDEX IF NOT EXISTS idx_documents_source_url ON documents(source_url)""",
|
||||
"""CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(
|
||||
title, text, source,
|
||||
content=documents,
|
||||
content_rowid=rowid
|
||||
)""",
|
||||
# Triggers to keep FTS in sync
|
||||
"""CREATE TRIGGER IF NOT EXISTS documents_ai AFTER INSERT ON documents BEGIN
|
||||
INSERT INTO documents_fts(rowid, title, text, source)
|
||||
VALUES (new.rowid, new.title, new.text, new.source);
|
||||
END""",
|
||||
"""CREATE TRIGGER IF NOT EXISTS documents_ad AFTER DELETE ON documents BEGIN
|
||||
INSERT INTO documents_fts(documents_fts, rowid, title, text, source)
|
||||
VALUES ('delete', old.rowid, old.title, old.text, old.source);
|
||||
END""",
|
||||
"""CREATE TRIGGER IF NOT EXISTS documents_au AFTER UPDATE ON documents BEGIN
|
||||
INSERT INTO documents_fts(documents_fts, rowid, title, text, source)
|
||||
VALUES ('delete', old.rowid, old.title, old.text, old.source);
|
||||
INSERT INTO documents_fts(rowid, title, text, source)
|
||||
VALUES (new.rowid, new.title, new.text, new.source);
|
||||
END""",
|
||||
"""CREATE TABLE IF NOT EXISTS extracted_tools (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
name TEXT NOT NULL DEFAULT '',
|
||||
aliases TEXT NOT NULL DEFAULT '[]',
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
capability TEXT NOT NULL DEFAULT '',
|
||||
target_platforms TEXT NOT NULL DEFAULT '[]',
|
||||
target_software TEXT NOT NULL DEFAULT '[]',
|
||||
cves TEXT NOT NULL DEFAULT '[]',
|
||||
source_doc_ids TEXT NOT NULL DEFAULT '[]',
|
||||
source_context TEXT NOT NULL DEFAULT '',
|
||||
mitre_techniques TEXT NOT NULL DEFAULT '[]',
|
||||
dedup_key TEXT NOT NULL DEFAULT ''
|
||||
)""",
|
||||
"""CREATE UNIQUE INDEX IF NOT EXISTS idx_tools_dedup ON extracted_tools(dedup_key)""",
|
||||
"""CREATE TABLE IF NOT EXISTS ttps (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
technique TEXT NOT NULL DEFAULT '',
|
||||
category TEXT NOT NULL DEFAULT '',
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
actor TEXT,
|
||||
mitre_id TEXT,
|
||||
source_doc_ids TEXT NOT NULL DEFAULT '[]',
|
||||
source_context TEXT NOT NULL DEFAULT '',
|
||||
dedup_key TEXT NOT NULL DEFAULT ''
|
||||
)""",
|
||||
"""CREATE UNIQUE INDEX IF NOT EXISTS idx_ttps_dedup ON ttps(dedup_key)""",
|
||||
"""CREATE TABLE IF NOT EXISTS tradecraft (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
method TEXT NOT NULL DEFAULT '',
|
||||
domain TEXT NOT NULL DEFAULT '',
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
operational_notes TEXT NOT NULL DEFAULT '',
|
||||
countermeasures TEXT NOT NULL DEFAULT '',
|
||||
source_doc_ids TEXT NOT NULL DEFAULT '[]',
|
||||
source_context TEXT NOT NULL DEFAULT ''
|
||||
)""",
|
||||
"""CREATE TABLE IF NOT EXISTS entities (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
name TEXT NOT NULL DEFAULT '',
|
||||
type TEXT NOT NULL DEFAULT '',
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
relationships TEXT NOT NULL DEFAULT '[]',
|
||||
source_doc_ids TEXT NOT NULL DEFAULT '[]',
|
||||
source_context TEXT NOT NULL DEFAULT ''
|
||||
)""",
|
||||
"""CREATE TABLE IF NOT EXISTS infrastructure (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
indicator TEXT NOT NULL DEFAULT '',
|
||||
indicator_type TEXT NOT NULL DEFAULT '',
|
||||
context TEXT NOT NULL DEFAULT '',
|
||||
source_doc_ids TEXT NOT NULL DEFAULT '[]',
|
||||
source_context TEXT NOT NULL DEFAULT '',
|
||||
collection TEXT NOT NULL DEFAULT '',
|
||||
doc_title TEXT NOT NULL DEFAULT '',
|
||||
doc_date TEXT
|
||||
)""",
|
||||
"""CREATE INDEX IF NOT EXISTS idx_infra_indicator ON infrastructure(indicator)""",
|
||||
"""CREATE INDEX IF NOT EXISTS idx_infra_type ON infrastructure(indicator_type)""",
|
||||
"""CREATE TABLE IF NOT EXISTS extraction_status (
|
||||
doc_id TEXT NOT NULL DEFAULT '',
|
||||
extractor TEXT NOT NULL DEFAULT '',
|
||||
status TEXT NOT NULL DEFAULT 'PENDING',
|
||||
error TEXT,
|
||||
last_run TEXT,
|
||||
token_count INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (doc_id, extractor)
|
||||
)""",
|
||||
"""CREATE TABLE IF NOT EXISTS token_usage (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
timestamp TEXT NOT NULL DEFAULT '',
|
||||
source TEXT NOT NULL DEFAULT '',
|
||||
doc_id TEXT NOT NULL DEFAULT '',
|
||||
extractor TEXT NOT NULL DEFAULT '',
|
||||
input_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
output_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
model TEXT NOT NULL DEFAULT ''
|
||||
)""",
|
||||
]
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class MigrationRunner:
|
||||
"""Run schema migrations on a SQLite database."""
|
||||
|
||||
def __init__(self, conn: sqlite3.Connection):
|
||||
self.conn = conn
|
||||
self._ensure_version_table()
|
||||
|
||||
def _ensure_version_table(self):
|
||||
self.conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS schema_version (
|
||||
version INTEGER PRIMARY KEY,
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
applied_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
)
|
||||
""")
|
||||
self.conn.commit()
|
||||
|
||||
def get_current_version(self) -> int:
|
||||
cursor = self.conn.execute(
|
||||
"SELECT MAX(version) FROM schema_version"
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return row[0] if row[0] is not None else 0
|
||||
|
||||
def run(self):
|
||||
current = self.get_current_version()
|
||||
pending = [m for m in MIGRATIONS if m['version'] > current]
|
||||
|
||||
if not pending:
|
||||
logger.debug("Database schema is up to date (version %d)", current)
|
||||
return
|
||||
|
||||
for migration in sorted(pending, key=lambda m: m['version']):
|
||||
logger.info(
|
||||
"Applying migration %d: %s",
|
||||
migration['version'],
|
||||
migration['description']
|
||||
)
|
||||
try:
|
||||
for sql in migration['sql']:
|
||||
self.conn.execute(sql)
|
||||
self.conn.execute(
|
||||
"INSERT INTO schema_version (version, description) VALUES (?, ?)",
|
||||
(migration['version'], migration['description'])
|
||||
)
|
||||
self.conn.commit()
|
||||
logger.info("Migration %d applied successfully", migration['version'])
|
||||
except Exception as e:
|
||||
self.conn.rollback()
|
||||
logger.error("Migration %d failed: %s", migration['version'], e)
|
||||
raise
|
||||
Reference in New Issue
Block a user