258 lines
8.9 KiB
Python
258 lines
8.9 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Resilient Redis → Supabase pgvector migration v2.
|
|
Fixes: smaller batches, retry on timeout, HNSW index dropped for speed.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
|
|
import httpx
|
|
import redis.asyncio as aioredis
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv("/app/.env", override=True)
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|
logger = logging.getLogger(__name__)
|
|
|
|
TARGET_DIM = 640
|
|
BATCH_SIZE = 50 # Smaller batches to avoid Supabase statement timeout
|
|
MAX_RETRIES = 3
|
|
|
|
COLLECTIONS = [
|
|
"wallet_profiles", # ~50K already migrated
|
|
"token_analysis",
|
|
"known_scams",
|
|
"scam_patterns",
|
|
"forensic_reports",
|
|
"market_intel",
|
|
"contract_audits",
|
|
"news_articles",
|
|
"transaction_patterns",
|
|
"general",
|
|
]
|
|
|
|
|
|
def _get_url():
|
|
return os.environ.get("SUPABASE_URL", "")
|
|
|
|
|
|
def _get_key():
|
|
return os.environ.get("SUPABASE_SERVICE_KEY", "") or os.environ.get("SUPABASE_SERVICE_ROLE_KEY", "")
|
|
|
|
|
|
def _get_headers():
|
|
key = _get_key()
|
|
return {
|
|
"apikey": key,
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
"Prefer": "resolution=merge-duplicates",
|
|
}
|
|
|
|
|
|
def pad_vector(vec: list, target_dim: int) -> list:
|
|
if len(vec) == target_dim:
|
|
return vec
|
|
if len(vec) > target_dim:
|
|
return vec[:target_dim]
|
|
return vec + [0.0] * (target_dim - len(vec))
|
|
|
|
|
|
async def get_existing_ids(client: httpx.AsyncClient) -> set:
|
|
"""Fetch existing doc IDs from rag_vectors for idempotency."""
|
|
existing = set()
|
|
offset = 0
|
|
limit = 1000
|
|
while True:
|
|
url = f"{_get_url()}/rest/v1/rag_vectors"
|
|
params = {"select": "id", "limit": str(limit), "offset": str(offset), "order": "id"}
|
|
try:
|
|
resp = await client.get(url, params=params, headers=_get_headers(), timeout=30)
|
|
if resp.status_code != 200:
|
|
logger.warning(f"Could not fetch existing IDs: {resp.status_code}")
|
|
break
|
|
rows = resp.json() if resp.text else []
|
|
if not rows:
|
|
break
|
|
for row in rows:
|
|
existing.add(row["id"])
|
|
if len(rows) < limit:
|
|
break
|
|
offset += limit
|
|
except Exception as e:
|
|
logger.warning(f"Error fetching existing IDs: {e}")
|
|
break
|
|
logger.info(f"Found {len(existing)} existing docs in rag_vectors")
|
|
return existing
|
|
|
|
|
|
async def insert_with_retry(client: httpx.AsyncClient, rows: list, batch_num: int) -> int:
|
|
"""Insert a batch with retry logic."""
|
|
url = f"{_get_url()}/rest/v1/rag_vectors"
|
|
params = {"on_conflict": "id"}
|
|
|
|
for attempt in range(MAX_RETRIES):
|
|
try:
|
|
resp = await client.post(url, json=rows, params=params, headers=_get_headers(), timeout=60)
|
|
if resp.status_code in (200, 201):
|
|
return len(rows)
|
|
elif resp.status_code == 500 and "timeout" in (resp.text or "").lower():
|
|
# Statement timeout — split into smaller batches
|
|
logger.warning(f"Batch {batch_num}: statement timeout (attempt {attempt + 1}), splitting")
|
|
if len(rows) > 10:
|
|
# Split into two halves
|
|
mid = len(rows) // 2
|
|
count = 0
|
|
for half in [rows[:mid], rows[mid:]]:
|
|
r2 = await client.post(url, json=half, params=params, headers=_get_headers(), timeout=90)
|
|
if r2.status_code in (200, 201):
|
|
count += len(half)
|
|
else:
|
|
# Last resort: insert one by one
|
|
for row in half:
|
|
try:
|
|
r3 = await client.post(
|
|
url,
|
|
json=[row],
|
|
params=params,
|
|
headers=_get_headers(),
|
|
timeout=30,
|
|
)
|
|
if r3.status_code in (200, 201):
|
|
count += 1
|
|
await asyncio.sleep(0.02)
|
|
except Exception:
|
|
pass
|
|
return count
|
|
else:
|
|
logger.error(f"Batch {batch_num} insert failed: {resp.status_code} {resp.text[:300]}")
|
|
return 0
|
|
except httpx.ReadTimeout:
|
|
logger.warning(f"Batch {batch_num}: timeout (attempt {attempt + 1})")
|
|
await asyncio.sleep(2**attempt)
|
|
except Exception as e:
|
|
logger.error(f"Batch {batch_num} exception: {e}")
|
|
await asyncio.sleep(1)
|
|
|
|
return 0
|
|
|
|
|
|
async def migrate_collection(r: aioredis.Redis, client: httpx.AsyncClient, collection: str, existing_ids: set) -> int:
|
|
"""Migrate all docs from one Redis collection to Supabase."""
|
|
idx_key = f"rag:idx:{collection}"
|
|
doc_ids = await r.smembers(idx_key)
|
|
if not doc_ids:
|
|
logger.info(f"Collection {collection}: empty, skipping")
|
|
return 0
|
|
|
|
total_docs = len(doc_ids)
|
|
to_migrate = [did for did in doc_ids if did not in existing_ids]
|
|
skipped = total_docs - len(to_migrate)
|
|
if skipped > 0:
|
|
logger.info(
|
|
f"Collection {collection}: {total_docs} total, skipping {skipped} already-migrated, {len(to_migrate)} to go"
|
|
)
|
|
|
|
migrated = 0
|
|
errors = 0
|
|
batch_rows = []
|
|
batch_count = 0
|
|
|
|
for _i, doc_id in enumerate(to_migrate):
|
|
doc_key = f"rag:{collection}:{doc_id}"
|
|
try:
|
|
data = await r.get(doc_key)
|
|
if not data:
|
|
continue
|
|
|
|
doc = json.loads(data)
|
|
vector = doc.get("vector", [])
|
|
content = doc.get("content", "") or ""
|
|
metadata = doc.get("metadata", {}) or {}
|
|
source = metadata.get("source", "") or doc.get("source", "") or ""
|
|
severity = metadata.get("severity", "") or doc.get("severity", "") or "medium"
|
|
chain = metadata.get("chain", "") or doc.get("chain", "") or ""
|
|
|
|
padded = pad_vector(vector, TARGET_DIM) if vector else [0.0] * TARGET_DIM
|
|
|
|
row = {
|
|
"id": str(doc_id),
|
|
"collection": collection,
|
|
"content": content[:10000],
|
|
"embedding": padded,
|
|
"metadata": json.dumps(metadata) if isinstance(metadata, dict) else str(metadata),
|
|
"source": source[:200] if source else "",
|
|
"severity": severity[:50] if severity else "medium",
|
|
"chain": chain[:50] if chain else "",
|
|
}
|
|
|
|
batch_rows.append(row)
|
|
|
|
if len(batch_rows) >= BATCH_SIZE:
|
|
count = await insert_with_retry(client, batch_rows, batch_count)
|
|
migrated += count
|
|
batch_count += 1
|
|
if batch_count % 10 == 0:
|
|
logger.info(f" Progress: {migrated}/{len(to_migrate)} migrated ({batch_count} batches)")
|
|
batch_rows = []
|
|
await asyncio.sleep(0.1)
|
|
|
|
except json.JSONDecodeError:
|
|
errors += 1
|
|
except Exception as e:
|
|
errors += 1
|
|
if errors <= 5:
|
|
logger.warning(f" Doc {doc_id}: {e}")
|
|
|
|
# Flush remaining
|
|
if batch_rows:
|
|
count = await insert_with_retry(client, batch_rows, batch_count)
|
|
migrated += count
|
|
|
|
logger.info(f"Collection {collection}: DONE {migrated}/{len(to_migrate)} ({errors} errors)")
|
|
return migrated
|
|
|
|
|
|
async def main():
|
|
logger.info(f"Starting migration v2 (batch_size={BATCH_SIZE})")
|
|
logger.info(f"Supabase: {_get_url()[:30]}...")
|
|
logger.info(f"Target dim: {TARGET_DIM}")
|
|
|
|
r = aioredis.Redis(
|
|
host=os.environ.get("REDIS_HOST", "rmi-redis"),
|
|
port=int(os.environ.get("REDIS_PORT", "6379")),
|
|
password=os.environ.get("REDIS_PASSWORD", ""),
|
|
decode_responses=True,
|
|
)
|
|
await r.ping()
|
|
logger.info("Redis connected")
|
|
|
|
async with httpx.AsyncClient(timeout=120) as client:
|
|
existing_ids = await get_existing_ids(client)
|
|
total = 0
|
|
start = time.time()
|
|
|
|
for coll in COLLECTIONS:
|
|
try:
|
|
n = await migrate_collection(r, client, coll, existing_ids)
|
|
total += n
|
|
except Exception as e:
|
|
logger.error(f"Collection {coll} FAILED: {e}")
|
|
|
|
elapsed = time.time() - start
|
|
logger.info("=" * 60)
|
|
logger.info(f"Migration v2 complete: {total} docs in {elapsed:.1f}s")
|
|
logger.info(
|
|
"Now rebuild HNSW index: CREATE INDEX idx_rag_vectors_hnsw ON rag_vectors USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 64);"
|
|
)
|
|
|
|
await r.aclose()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|