Spaces:
Running
Running
File size: 4,456 Bytes
1e3eb08 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | #!/usr/bin/env python3
"""
One-time migration: export all kv_store rows from local Dolt → agentmemory.db (SQLite).
Run ONCE while local Dolt SQL server is running on 127.0.0.1:3306.
"""
import os, sys, json, sqlite3, time
env_path = os.path.expanduser("~/.agentmemory/.env")
with open(env_path) as f:
for line in f:
line = line.strip()
if line and not line.startswith("#") and "=" in line:
k, v = line.split("=", 1)
os.environ[k.strip()] = v.strip().strip('"').strip("'")
try:
import pymysql
except ImportError:
print("[migrate] pymysql not installed. Run: pip install pymysql")
sys.exit(1)
DOLT_HOST = os.environ.get("DOLT_HOST", "127.0.0.1")
DOLT_PORT = int(os.environ.get("DOLT_PORT", "3306"))
DOLT_USER = os.environ.get("DOLT_USER", "root")
DOLT_PASS = os.environ.get("DOLT_PASSWORD", "")
DOLT_DB = os.environ.get("DOLT_DATABASE", "agentmemory")
SQLITE_PATH = os.path.expanduser("~/.agentmemory/agentmemory.db")
print(f"[migrate] Connecting to Dolt at {DOLT_HOST}:{DOLT_PORT}/{DOLT_DB}...")
try:
dolt = pymysql.connect(
host=DOLT_HOST, port=DOLT_PORT, user=DOLT_USER,
password=DOLT_PASS, database=DOLT_DB,
charset="utf8mb4", cursorclass=pymysql.cursors.DictCursor, autocommit=True
)
except Exception as e:
print(f"[migrate] Dolt connection failed: {e}")
print("[migrate] Make sure Dolt SQL server is running: dolt sql-server --host 127.0.0.1 --port 3306 --data-dir ~/.agentmemory/dolt")
sys.exit(1)
print(f"[migrate] Opening SQLite at {SQLITE_PATH}...")
db = sqlite3.connect(SQLITE_PATH)
db.execute("PRAGMA journal_mode=WAL")
db.execute("""
CREATE TABLE IF NOT EXISTS kv_store (
scope TEXT NOT NULL,
key TEXT NOT NULL,
value TEXT NOT NULL,
PRIMARY KEY (scope, key)
)
""")
db.execute("CREATE INDEX IF NOT EXISTS idx_kv_scope ON kv_store(scope)")
db.execute("""
CREATE TABLE IF NOT EXISTS audit_log (
id INTEGER PRIMARY KEY AUTOINCREMENT,
ts INTEGER NOT NULL,
agent_id TEXT NOT NULL,
message TEXT NOT NULL
)
""")
db.commit()
# Read all rows from Dolt
with dolt.cursor() as cur:
cur.execute("SELECT COUNT(*) AS n FROM kv_store")
total = cur.fetchone()["n"]
print(f"[migrate] Found {total} rows in Dolt kv_store")
cur.execute("SELECT scope, `key`, value FROM kv_store")
rows = cur.fetchall()
inserted = 0
for row in rows:
try:
db.execute(
"INSERT OR REPLACE INTO kv_store (scope, key, value) VALUES (?, ?, ?)",
(row["scope"], row["key"], row["value"])
)
inserted += 1
except Exception as e:
print(f"[migrate] row error ({row['scope']}/{row['key']}): {e}")
db.commit()
print(f"[migrate] Inserted {inserted}/{total} rows into SQLite")
# Try to migrate dolt_log as audit entries
try:
with dolt.cursor() as cur:
cur.execute("SELECT committer, date, message FROM dolt_log ORDER BY date ASC LIMIT 1000")
logs = cur.fetchall()
for l in logs:
ts = int(l["date"].timestamp() * 1000) if hasattr(l["date"], "timestamp") else int(time.time() * 1000)
db.execute(
"INSERT INTO audit_log (ts, agent_id, message) VALUES (?, ?, ?)",
(ts, l["committer"] or "unknown", l["message"] or "")
)
db.commit()
print(f"[migrate] Migrated {len(logs)} dolt_log entries to audit_log")
except Exception as e:
print(f"[migrate] dolt_log migration skipped: {e}")
dolt.close()
db.close()
size_mb = os.path.getsize(SQLITE_PATH) / 1024 / 1024
print(f"[migrate] Done! SQLite file: {SQLITE_PATH} ({size_mb:.1f} MB)")
print("[migrate] Now run: python push_local_data.py (or python -c below)")
print("""
python -c "
import os, sys
env = os.path.expanduser('~/.agentmemory/.env')
with open(env) as f:
for l in f:
l=l.strip()
if l and not l.startswith('#') and '=' in l:
k,v=l.split('=',1); os.environ[k.strip()]=v.strip().strip('\\\"').strip(\\\"'\\\")
from huggingface_hub import HfApi
api = HfApi(token=os.environ['HF_TOKEN'])
api.upload_file(
path_or_fileobj=os.path.expanduser('~/.agentmemory/agentmemory.db'),
path_in_repo='agentmemory.db',
repo_id=os.environ.get('AGENTMEMORY_DATASET_REPO','Yash030/agentmemory-python-data'),
repo_type='dataset', token=os.environ['HF_TOKEN'],
commit_message='restore: upload migrated SQLite DB'
)
print('Uploaded!')
"
""")
|