Initial commit: Egregore db service
PostgreSQL message storage API with asyncpg connection pooling and full-text search. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
commit
291d664051
5 changed files with 686 additions and 0 deletions
45
.gitignore
vendored
Normal file
45
.gitignore
vendored
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
.Python
|
||||
*.egg-info/
|
||||
.eggs/
|
||||
dist/
|
||||
build/
|
||||
|
||||
# Environment
|
||||
.env
|
||||
.env.*
|
||||
*.local
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
|
||||
# Secrets - BE PARANOID
|
||||
*.pem
|
||||
*.key
|
||||
*.crt
|
||||
*.p12
|
||||
credentials*
|
||||
secrets*
|
||||
tokens*
|
||||
*_secret*
|
||||
*_token*
|
||||
*.credentials
|
||||
|
||||
# Logs and data
|
||||
*.log
|
||||
*.db
|
||||
*.sqlite
|
||||
*.sqlite3
|
||||
*.backup
|
||||
chat.db*
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
28
__init__.py
Normal file
28
__init__.py
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
"""
|
||||
Egregore Database - Message storage and retrieval
|
||||
|
||||
This module handles all database operations for the chat system.
|
||||
Currently SQLite, designed for easy migration to PostgreSQL later.
|
||||
"""
|
||||
|
||||
from .messages import (
|
||||
init_db,
|
||||
save_message,
|
||||
save_response_blocks,
|
||||
get_messages,
|
||||
get_conversation_history,
|
||||
search_messages,
|
||||
MESSAGE_PRIORITIES,
|
||||
get_priority_for_type,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"init_db",
|
||||
"save_message",
|
||||
"save_response_blocks",
|
||||
"get_messages",
|
||||
"get_conversation_history",
|
||||
"search_messages",
|
||||
"MESSAGE_PRIORITIES",
|
||||
"get_priority_for_type",
|
||||
]
|
||||
116
main.py
Normal file
116
main.py
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Egregore DB Service - Message storage API
|
||||
|
||||
Provides HTTP API for message storage and retrieval.
|
||||
Runs on port 8082.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import FastAPI, Query
|
||||
from pydantic import BaseModel
|
||||
|
||||
from messages import (
|
||||
init_db,
|
||||
save_message,
|
||||
save_response_blocks,
|
||||
get_messages,
|
||||
get_conversation_history,
|
||||
search_messages,
|
||||
close_pool,
|
||||
)
|
||||
|
||||
app = FastAPI(title="Egregore DB Service", docs_url="/docs")
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
async def startup():
|
||||
await init_db()
|
||||
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def shutdown():
|
||||
await close_pool()
|
||||
|
||||
|
||||
# Request models
|
||||
class SaveMessageRequest(BaseModel):
|
||||
role: str
|
||||
content: str
|
||||
msg_type: str = "text"
|
||||
group_id: Optional[str] = None
|
||||
metadata: Optional[dict] = None
|
||||
priority: Optional[int] = None
|
||||
|
||||
|
||||
class SaveBlocksRequest(BaseModel):
|
||||
blocks: list
|
||||
group_id: str
|
||||
|
||||
|
||||
# Endpoints
|
||||
@app.post("/messages")
|
||||
async def api_save_message(req: SaveMessageRequest):
|
||||
"""Save a single message"""
|
||||
msg_id = await save_message(
|
||||
role=req.role,
|
||||
content=req.content,
|
||||
msg_type=req.msg_type,
|
||||
group_id=req.group_id,
|
||||
metadata=req.metadata,
|
||||
priority=req.priority
|
||||
)
|
||||
return {"id": msg_id}
|
||||
|
||||
|
||||
@app.post("/messages/blocks")
|
||||
async def api_save_blocks(req: SaveBlocksRequest):
|
||||
"""Save multiple response blocks"""
|
||||
saved = await save_response_blocks(req.blocks, req.group_id)
|
||||
return {"messages": saved}
|
||||
|
||||
|
||||
@app.get("/messages")
|
||||
async def api_get_messages(
|
||||
limit: int = 50,
|
||||
before: Optional[int] = None,
|
||||
msg_type: Optional[str] = Query(None, alias="type")
|
||||
):
|
||||
"""Get messages with pagination"""
|
||||
messages, has_more = await get_messages(
|
||||
limit=min(limit, 100),
|
||||
before_id=before,
|
||||
msg_type=msg_type
|
||||
)
|
||||
return {"messages": messages, "has_more": has_more}
|
||||
|
||||
|
||||
@app.get("/messages/history")
|
||||
async def api_get_history(limit: int = 100):
|
||||
"""Get conversation history in Claude API format"""
|
||||
history = await get_conversation_history(limit=limit)
|
||||
return {"history": history}
|
||||
|
||||
|
||||
@app.get("/messages/search")
|
||||
async def api_search(
|
||||
q: str,
|
||||
limit: int = 20,
|
||||
msg_type: Optional[str] = Query(None, alias="type")
|
||||
):
|
||||
"""Search messages by content"""
|
||||
results = await search_messages(q, limit=limit, msg_type=msg_type)
|
||||
return {"results": results}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
"""Health check endpoint"""
|
||||
return {"status": "ok", "service": "db"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="127.0.0.1", port=8082)
|
||||
371
messages.py
Normal file
371
messages.py
Normal file
|
|
@ -0,0 +1,371 @@
|
|||
"""
|
||||
Egregore Database - Message storage operations (PostgreSQL)
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
import asyncpg
|
||||
|
||||
# Database connection URL - can be overridden via environment
|
||||
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://egregore:egregore_db_pass@localhost/egregore")
|
||||
|
||||
# Connection pool
|
||||
_pool: Optional[asyncpg.Pool] = None
|
||||
|
||||
|
||||
def set_db_url(url: str):
|
||||
"""Set the database URL (call before init_db)"""
|
||||
global DATABASE_URL
|
||||
DATABASE_URL = url
|
||||
|
||||
|
||||
async def get_pool() -> asyncpg.Pool:
|
||||
"""Get or create the connection pool"""
|
||||
global _pool
|
||||
if _pool is None:
|
||||
_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
|
||||
return _pool
|
||||
|
||||
|
||||
async def close_pool():
|
||||
"""Close the connection pool"""
|
||||
global _pool
|
||||
if _pool:
|
||||
await _pool.close()
|
||||
_pool = None
|
||||
|
||||
|
||||
# Message type priorities for notifications
|
||||
MESSAGE_PRIORITIES = {
|
||||
"text": 2, # Regular messages - notify
|
||||
"tool_use": 0, # Tool invocation - no notify
|
||||
"tool_result": 0, # Tool output - no notify
|
||||
"question": 3, # Questions to user - urgent notify
|
||||
"mode_change": 1, # State transitions - silent
|
||||
"thinking": 0, # Reasoning process - no notify
|
||||
"error": 2, # Error messages - notify
|
||||
}
|
||||
|
||||
|
||||
def get_priority_for_type(msg_type: str, content: str = "") -> int:
|
||||
"""Get priority for a message type, with question detection"""
|
||||
base_priority = MESSAGE_PRIORITIES.get(msg_type, 0)
|
||||
if msg_type == "text" and content.strip().endswith("?"):
|
||||
return MESSAGE_PRIORITIES["question"]
|
||||
return base_priority
|
||||
|
||||
|
||||
async def init_db():
|
||||
"""Initialize PostgreSQL database with messages table"""
|
||||
pool = await get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id SERIAL PRIMARY KEY,
|
||||
role TEXT NOT NULL,
|
||||
type TEXT NOT NULL DEFAULT 'text',
|
||||
content TEXT NOT NULL,
|
||||
group_id TEXT,
|
||||
metadata JSONB,
|
||||
priority INTEGER DEFAULT 0,
|
||||
timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
)
|
||||
""")
|
||||
|
||||
# Create indexes for efficient querying
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_messages_role ON messages(role)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_messages_type ON messages(type)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_messages_group_id ON messages(group_id)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_messages_priority ON messages(priority)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_messages_timestamp ON messages(timestamp)")
|
||||
|
||||
# Full-text search index on content
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_content_search
|
||||
ON messages USING gin(to_tsvector('english', content))
|
||||
""")
|
||||
|
||||
|
||||
async def save_message(
|
||||
role: str,
|
||||
content: str,
|
||||
msg_type: str = "text",
|
||||
group_id: Optional[str] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
priority: Optional[int] = None
|
||||
) -> int:
|
||||
"""Save a single message row"""
|
||||
pool = await get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
if priority is None:
|
||||
priority = get_priority_for_type(msg_type, content)
|
||||
|
||||
row = await conn.fetchrow(
|
||||
"""INSERT INTO messages (role, type, content, group_id, metadata, priority)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
RETURNING id""",
|
||||
role, msg_type, content, group_id, json.dumps(metadata) if metadata else None, priority
|
||||
)
|
||||
return row['id']
|
||||
|
||||
|
||||
async def save_response_blocks(blocks: list, group_id: str) -> list:
|
||||
"""
|
||||
Save each response block as a separate row.
|
||||
Returns list of saved message dicts with IDs for frontend.
|
||||
"""
|
||||
saved_messages = []
|
||||
pool = await get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
timestamp = datetime.utcnow()
|
||||
|
||||
for block in blocks:
|
||||
block_type = block.get("type", "text")
|
||||
content = ""
|
||||
metadata = None
|
||||
priority = MESSAGE_PRIORITIES.get(block_type, 0)
|
||||
|
||||
if block_type == "text":
|
||||
content = block.get("content", "")
|
||||
if content.strip().endswith("?"):
|
||||
priority = MESSAGE_PRIORITIES["question"]
|
||||
elif block_type == "tool_use":
|
||||
content = json.dumps(block.get("input", {}))
|
||||
metadata = {"tool_name": block.get("name"), "tool_id": block.get("id")}
|
||||
elif block_type == "tool_result":
|
||||
content = block.get("content", "")
|
||||
metadata = {"tool_name": block.get("tool_name"), "tool_use_id": block.get("tool_use_id")}
|
||||
else:
|
||||
content = block.get("content", "")
|
||||
|
||||
row = await conn.fetchrow(
|
||||
"""INSERT INTO messages (role, type, content, group_id, metadata, priority, timestamp)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
RETURNING id""",
|
||||
"assistant", block_type, content, group_id,
|
||||
json.dumps(metadata) if metadata else None, priority, timestamp
|
||||
)
|
||||
msg_id = row['id']
|
||||
|
||||
saved_messages.append({
|
||||
"id": msg_id,
|
||||
"role": "assistant",
|
||||
"type": block_type,
|
||||
"content": content,
|
||||
"group_id": group_id,
|
||||
"metadata": metadata,
|
||||
"priority": priority,
|
||||
"timestamp": timestamp.isoformat()
|
||||
})
|
||||
|
||||
return saved_messages
|
||||
|
||||
|
||||
async def get_messages(
|
||||
limit: int = 50,
|
||||
before_id: int = None,
|
||||
msg_type: str = None
|
||||
) -> tuple[list[dict], bool]:
|
||||
"""Get messages with pagination. Returns (messages, has_more)"""
|
||||
pool = await get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
params = []
|
||||
where_clauses = []
|
||||
param_idx = 1
|
||||
|
||||
if before_id:
|
||||
where_clauses.append(f"id < ${param_idx}")
|
||||
params.append(before_id)
|
||||
param_idx += 1
|
||||
if msg_type:
|
||||
where_clauses.append(f"type = ${param_idx}")
|
||||
params.append(msg_type)
|
||||
param_idx += 1
|
||||
|
||||
where_sql = " AND ".join(where_clauses) if where_clauses else "TRUE"
|
||||
params.append(limit + 1)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""SELECT id, role, type, content, group_id, metadata, priority, timestamp
|
||||
FROM messages
|
||||
WHERE {where_sql}
|
||||
ORDER BY id DESC
|
||||
LIMIT ${param_idx}""",
|
||||
*params
|
||||
)
|
||||
|
||||
has_more = len(rows) > limit
|
||||
rows = rows[:limit]
|
||||
|
||||
messages = []
|
||||
for row in rows:
|
||||
metadata = None
|
||||
if row['metadata']:
|
||||
try:
|
||||
metadata = json.loads(row['metadata']) if isinstance(row['metadata'], str) else row['metadata']
|
||||
except:
|
||||
pass
|
||||
messages.append({
|
||||
"id": row['id'],
|
||||
"role": row['role'],
|
||||
"type": row['type'],
|
||||
"content": row['content'],
|
||||
"group_id": row['group_id'],
|
||||
"metadata": metadata,
|
||||
"priority": row['priority'],
|
||||
"timestamp": row['timestamp'].isoformat() if row['timestamp'] else None
|
||||
})
|
||||
|
||||
return list(reversed(messages)), has_more
|
||||
|
||||
|
||||
async def get_conversation_history(limit: int = 100) -> list[dict]:
|
||||
"""
|
||||
Reconstruct Claude API message format from individual rows.
|
||||
Groups assistant messages by group_id to build proper content arrays.
|
||||
"""
|
||||
messages, _ = await get_messages(limit)
|
||||
|
||||
api_messages = []
|
||||
current_group = None
|
||||
current_assistant_content = []
|
||||
|
||||
for msg in messages:
|
||||
if msg["role"] == "user":
|
||||
if current_assistant_content:
|
||||
api_messages.append({
|
||||
"role": "assistant",
|
||||
"content": current_assistant_content
|
||||
})
|
||||
current_assistant_content = []
|
||||
current_group = None
|
||||
|
||||
api_messages.append({
|
||||
"role": "user",
|
||||
"content": msg["content"]
|
||||
})
|
||||
elif msg["role"] == "assistant":
|
||||
if msg["group_id"] != current_group:
|
||||
if current_assistant_content:
|
||||
api_messages.append({
|
||||
"role": "assistant",
|
||||
"content": current_assistant_content
|
||||
})
|
||||
current_assistant_content = []
|
||||
current_group = msg["group_id"]
|
||||
|
||||
if msg["type"] == "text":
|
||||
current_assistant_content.append({
|
||||
"type": "text",
|
||||
"text": msg["content"]
|
||||
})
|
||||
elif msg["type"] == "tool_use":
|
||||
tool_input = {}
|
||||
try:
|
||||
tool_input = json.loads(msg["content"])
|
||||
except:
|
||||
pass
|
||||
metadata = msg.get("metadata") or {}
|
||||
current_assistant_content.append({
|
||||
"type": "tool_use",
|
||||
"id": metadata.get("tool_id", str(uuid.uuid4())),
|
||||
"name": metadata.get("tool_name", "unknown"),
|
||||
"input": tool_input
|
||||
})
|
||||
elif msg["type"] == "tool_result":
|
||||
if current_assistant_content:
|
||||
api_messages.append({
|
||||
"role": "assistant",
|
||||
"content": current_assistant_content
|
||||
})
|
||||
current_assistant_content = []
|
||||
metadata = msg.get("metadata") or {}
|
||||
api_messages.append({
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": metadata.get("tool_use_id", ""),
|
||||
"content": msg["content"]
|
||||
}]
|
||||
})
|
||||
|
||||
if current_assistant_content:
|
||||
api_messages.append({
|
||||
"role": "assistant",
|
||||
"content": current_assistant_content
|
||||
})
|
||||
|
||||
return api_messages
|
||||
|
||||
|
||||
async def search_messages(
|
||||
query: str,
|
||||
limit: int = 20,
|
||||
msg_type: str = None
|
||||
) -> list[dict]:
|
||||
"""Search messages using PostgreSQL full-text search"""
|
||||
if len(query) < 2:
|
||||
return []
|
||||
|
||||
pool = await get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
params = [query]
|
||||
param_idx = 2
|
||||
type_filter = ""
|
||||
|
||||
if msg_type:
|
||||
type_filter = f"AND type = ${param_idx}"
|
||||
params.append(msg_type)
|
||||
param_idx += 1
|
||||
|
||||
params.append(min(limit, 50))
|
||||
|
||||
# Use PostgreSQL full-text search with fallback to ILIKE
|
||||
rows = await conn.fetch(
|
||||
f"""SELECT id, role, type, content, group_id, metadata, priority, timestamp,
|
||||
ts_headline('english', content, plainto_tsquery('english', $1),
|
||||
'StartSel=**, StopSel=**, MaxWords=50, MinWords=20') as snippet
|
||||
FROM messages
|
||||
WHERE (to_tsvector('english', content) @@ plainto_tsquery('english', $1)
|
||||
OR content ILIKE '%' || $1 || '%')
|
||||
{type_filter}
|
||||
ORDER BY id DESC
|
||||
LIMIT ${param_idx}""",
|
||||
*params
|
||||
)
|
||||
|
||||
results = []
|
||||
for row in rows:
|
||||
metadata = None
|
||||
if row['metadata']:
|
||||
try:
|
||||
metadata = json.loads(row['metadata']) if isinstance(row['metadata'], str) else row['metadata']
|
||||
except:
|
||||
pass
|
||||
|
||||
# Use ts_headline snippet, fallback to manual snippet
|
||||
snippet = row['snippet'] if row['snippet'] else row['content'][:100]
|
||||
|
||||
results.append({
|
||||
"id": row['id'],
|
||||
"role": row['role'],
|
||||
"type": row['type'],
|
||||
"content": row['content'],
|
||||
"group_id": row['group_id'],
|
||||
"metadata": metadata,
|
||||
"priority": row['priority'],
|
||||
"timestamp": row['timestamp'].isoformat() if row['timestamp'] else None,
|
||||
"snippet": snippet
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# Legacy compatibility - for migration
|
||||
def set_db_path(path: str):
|
||||
"""Legacy function for SQLite compatibility - ignored for PostgreSQL"""
|
||||
pass
|
||||
126
migrate_to_postgres.py
Normal file
126
migrate_to_postgres.py
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Migration script: SQLite to PostgreSQL
|
||||
|
||||
Migrates data from the old SQLite database to PostgreSQL.
|
||||
Run once after setting up PostgreSQL.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
from datetime import datetime
|
||||
|
||||
import asyncpg
|
||||
|
||||
# Paths
|
||||
SQLITE_DB = "/home/admin/services/db/chat.db"
|
||||
POSTGRES_URL = os.getenv("DATABASE_URL", "postgresql://egregore:egregore_db_pass@localhost/egregore")
|
||||
|
||||
|
||||
async def migrate():
|
||||
print(f"Connecting to PostgreSQL...")
|
||||
pool = await asyncpg.create_pool(POSTGRES_URL)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
# Create table if not exists
|
||||
print("Creating messages table...")
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id SERIAL PRIMARY KEY,
|
||||
role TEXT NOT NULL,
|
||||
type TEXT NOT NULL DEFAULT 'text',
|
||||
content TEXT NOT NULL,
|
||||
group_id TEXT,
|
||||
metadata JSONB,
|
||||
priority INTEGER DEFAULT 0,
|
||||
timestamp TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
)
|
||||
""")
|
||||
|
||||
# Create indexes
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_messages_role ON messages(role)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_messages_type ON messages(type)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_messages_group_id ON messages(group_id)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_messages_priority ON messages(priority)")
|
||||
await conn.execute("CREATE INDEX IF NOT EXISTS idx_messages_timestamp ON messages(timestamp)")
|
||||
await conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_content_search
|
||||
ON messages USING gin(to_tsvector('english', content))
|
||||
""")
|
||||
|
||||
# Check if SQLite database exists
|
||||
if not os.path.exists(SQLITE_DB):
|
||||
print(f"No SQLite database found at {SQLITE_DB}, nothing to migrate.")
|
||||
await pool.close()
|
||||
return
|
||||
|
||||
# Connect to SQLite
|
||||
print(f"Reading from SQLite: {SQLITE_DB}")
|
||||
sqlite_conn = sqlite3.connect(SQLITE_DB)
|
||||
sqlite_conn.row_factory = sqlite3.Row
|
||||
cursor = sqlite_conn.cursor()
|
||||
|
||||
# Check if messages_v2 table exists
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='messages_v2'")
|
||||
if not cursor.fetchone():
|
||||
print("No messages_v2 table in SQLite, nothing to migrate.")
|
||||
sqlite_conn.close()
|
||||
await pool.close()
|
||||
return
|
||||
|
||||
# Fetch all messages from SQLite
|
||||
cursor.execute("""
|
||||
SELECT id, role, type, content, group_id, metadata, priority, timestamp
|
||||
FROM messages_v2
|
||||
ORDER BY id ASC
|
||||
""")
|
||||
rows = cursor.fetchall()
|
||||
print(f"Found {len(rows)} messages to migrate")
|
||||
|
||||
if not rows:
|
||||
print("No messages to migrate.")
|
||||
sqlite_conn.close()
|
||||
await pool.close()
|
||||
return
|
||||
|
||||
# Insert into PostgreSQL
|
||||
async with pool.acquire() as conn:
|
||||
migrated = 0
|
||||
for row in rows:
|
||||
timestamp = datetime.fromisoformat(row['timestamp']) if row['timestamp'] else datetime.utcnow()
|
||||
|
||||
await conn.execute(
|
||||
"""INSERT INTO messages (role, type, content, group_id, metadata, priority, timestamp)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)""",
|
||||
row['role'],
|
||||
row['type'],
|
||||
row['content'],
|
||||
row['group_id'],
|
||||
row['metadata'], # Already JSON string, PostgreSQL will handle it
|
||||
row['priority'],
|
||||
timestamp
|
||||
)
|
||||
migrated += 1
|
||||
|
||||
print(f"Migrated {migrated} messages to PostgreSQL")
|
||||
|
||||
# Verify
|
||||
async with pool.acquire() as conn:
|
||||
count = await conn.fetchval("SELECT COUNT(*) FROM messages")
|
||||
print(f"PostgreSQL now has {count} messages")
|
||||
|
||||
sqlite_conn.close()
|
||||
await pool.close()
|
||||
|
||||
# Backup SQLite database
|
||||
backup_path = SQLITE_DB + ".backup"
|
||||
print(f"Backing up SQLite database to {backup_path}")
|
||||
os.rename(SQLITE_DB, backup_path)
|
||||
|
||||
print("Migration complete!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(migrate())
|
||||
Loading…
Add table
Add a link
Reference in a new issue