agnes-the-ai-analyst/scripts/migrate_json_to_duckdb.py

182 lines
6.4 KiB
Python

"""
One-time migration: JSON files -> DuckDB.
Usage: python -m scripts.migrate_json_to_duckdb [--data-dir /data]
Idempotent -- safe to run multiple times. Uses UPSERT to avoid duplicates.
"""
import json
import logging
import os
import sys
from pathlib import Path
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger(__name__)
def _load_json(path: str):
try:
with open(path) as f:
return json.load(f)
except (FileNotFoundError, json.JSONDecodeError) as e:
logger.warning(f"Skipping {path}: {e}")
return None
def migrate_all(data_dir: str = None) -> dict:
if data_dir:
os.environ["DATA_DIR"] = data_dir
data = Path(data_dir or os.environ.get("DATA_DIR", "./data"))
from src.db import get_system_db
from src.repositories.sync_state import SyncStateRepository
from src.repositories.knowledge import KnowledgeRepository
from src.repositories.notifications import TelegramRepository, PendingCodeRepository
from src.repositories.users import UserRepository
from src.repositories.table_registry import TableRegistryRepository
from src.repositories.profiles import ProfileRepository
conn = get_system_db()
stats = {}
# 1. Sync state
sync_data = _load_json(str(data / "src_data" / "metadata" / "sync_state.json"))
count = 0
if sync_data and isinstance(sync_data, dict):
repo = SyncStateRepository(conn)
tables = sync_data.get("tables", sync_data)
if not isinstance(tables, dict):
tables = {}
for table_id, info in tables.items():
if isinstance(info, dict):
repo.update_sync(
table_id=table_id,
rows=info.get("rows", 0),
file_size_bytes=info.get("file_size_bytes", 0),
hash=info.get("hash", ""),
uncompressed_size_bytes=info.get("uncompressed_size_bytes", 0),
columns=info.get("columns", 0),
)
count += 1
stats["sync_state"] = count
logger.info(f"Migrated {count} sync state entries")
# 2. Knowledge items
knowledge = _load_json(str(data / "corporate-memory" / "knowledge.json"))
count = 0
if knowledge and isinstance(knowledge, list):
repo = KnowledgeRepository(conn)
for item in knowledge:
if not isinstance(item, dict):
continue
item_id = item.get("id", "")
if not item_id:
continue
# Check if exists (idempotent)
existing = repo.get_by_id(item_id)
if existing:
continue
repo.create(
id=item_id,
title=item.get("title", ""),
content=item.get("content", ""),
category=item.get("category", ""),
source_user=item.get("source_user"),
tags=item.get("tags"),
status=item.get("status", "pending"),
)
count += 1
stats["knowledge"] = count
logger.info(f"Migrated {count} knowledge items")
# 3. Telegram users
telegram = _load_json(str(data / "notifications" / "telegram_users.json"))
count = 0
if telegram and isinstance(telegram, dict):
repo = TelegramRepository(conn)
for user_id, info in telegram.items():
if isinstance(info, dict) and "chat_id" in info:
repo.link_user(user_id, chat_id=info["chat_id"])
count += 1
stats["telegram"] = count
logger.info(f"Migrated {count} telegram links")
# 4. Password users
password_users = _load_json(str(data / "auth" / "password_users.json"))
count = 0
if password_users and isinstance(password_users, dict):
user_repo = UserRepository(conn)
for email, info in password_users.items():
if not isinstance(info, dict):
continue
existing = user_repo.get_by_email(email)
if existing:
continue
import uuid
user_repo.create(
id=str(uuid.uuid4()),
email=email,
name=info.get("name", email.split("@")[0]),
role="analyst",
password_hash=info.get("password_hash"),
)
count += 1
stats["users"] = count
logger.info(f"Migrated {count} password users")
# 5. Table registry
registry = _load_json(str(data / "src_data" / "metadata" / "table_registry.json"))
count = 0
if registry and isinstance(registry, dict):
repo = TableRegistryRepository(conn)
tables_list = registry.get("tables", [])
if isinstance(tables_list, list):
for t in tables_list:
if not isinstance(t, dict):
continue
tid = t.get("id", "")
if not tid:
continue
existing = repo.get(tid)
if existing:
continue
repo.register(
id=tid,
name=t.get("name", tid),
folder=t.get("folder"),
sync_strategy=t.get("sync_strategy"),
primary_key=t.get("primary_key"),
description=t.get("description"),
registered_by=t.get("registered_by"),
)
count += 1
stats["table_registry"] = count
logger.info(f"Migrated {count} table registry entries")
# 6. Profiles
profiles = _load_json(str(data / "src_data" / "metadata" / "profiles.json"))
count = 0
if profiles and isinstance(profiles, dict):
repo = ProfileRepository(conn)
tables_data = profiles.get("tables", profiles)
if isinstance(tables_data, dict):
for table_id, profile in tables_data.items():
if isinstance(profile, dict):
repo.save(table_id, profile)
count += 1
stats["profiles"] = count
logger.info(f"Migrated {count} table profiles")
conn.close()
logger.info(f"Migration complete: {stats}")
return stats
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Migrate JSON state to DuckDB")
parser.add_argument("--data-dir", default=None)
args = parser.parse_args()
migrate_all(args.data_dir)