|
|
"""
|
|
|
D1 Database integration for OpenManus
|
|
|
Provides interface to Cloudflare D1 database operations
|
|
|
"""
|
|
|
|
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
|
|
from app.logger import logger
|
|
|
|
|
|
from .client import CloudflareClient, CloudflareError
|
|
|
|
|
|
|
|
|
class D1Database:
|
|
|
"""Cloudflare D1 Database client"""
|
|
|
|
|
|
def __init__(self, client: CloudflareClient, database_id: str):
|
|
|
self.client = client
|
|
|
self.database_id = database_id
|
|
|
self.base_endpoint = f"accounts/{client.account_id}/d1/database/{database_id}"
|
|
|
|
|
|
async def execute_query(
|
|
|
self, sql: str, params: Optional[List[Any]] = None, use_worker: bool = True
|
|
|
) -> Dict[str, Any]:
|
|
|
"""Execute a SQL query"""
|
|
|
|
|
|
query_data = {"sql": sql}
|
|
|
|
|
|
if params:
|
|
|
query_data["params"] = params
|
|
|
|
|
|
try:
|
|
|
if use_worker:
|
|
|
|
|
|
response = await self.client.post(
|
|
|
"api/database/query", data=query_data, use_worker=True
|
|
|
)
|
|
|
else:
|
|
|
|
|
|
response = await self.client.post(
|
|
|
f"{self.base_endpoint}/query", data=query_data
|
|
|
)
|
|
|
|
|
|
return response
|
|
|
|
|
|
except CloudflareError as e:
|
|
|
logger.error(f"D1 query execution failed: {e}")
|
|
|
raise
|
|
|
|
|
|
async def batch_execute(
|
|
|
self, queries: List[Dict[str, Any]], use_worker: bool = True
|
|
|
) -> Dict[str, Any]:
|
|
|
"""Execute multiple queries in a batch"""
|
|
|
|
|
|
batch_data = {"queries": queries}
|
|
|
|
|
|
try:
|
|
|
if use_worker:
|
|
|
response = await self.client.post(
|
|
|
"api/database/batch", data=batch_data, use_worker=True
|
|
|
)
|
|
|
else:
|
|
|
response = await self.client.post(
|
|
|
f"{self.base_endpoint}/query", data=batch_data
|
|
|
)
|
|
|
|
|
|
return response
|
|
|
|
|
|
except CloudflareError as e:
|
|
|
logger.error(f"D1 batch execution failed: {e}")
|
|
|
raise
|
|
|
|
|
|
|
|
|
async def create_user(
|
|
|
self,
|
|
|
user_id: str,
|
|
|
username: str,
|
|
|
email: Optional[str] = None,
|
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
|
) -> Dict[str, Any]:
|
|
|
"""Create a new user"""
|
|
|
|
|
|
sql = """
|
|
|
INSERT INTO users (id, username, email, metadata)
|
|
|
VALUES (?, ?, ?, ?)
|
|
|
ON CONFLICT(id) DO UPDATE SET
|
|
|
username = excluded.username,
|
|
|
email = excluded.email,
|
|
|
metadata = excluded.metadata,
|
|
|
updated_at = strftime('%s', 'now')
|
|
|
"""
|
|
|
|
|
|
import json
|
|
|
|
|
|
params = [user_id, username, email, json.dumps(metadata or {})]
|
|
|
|
|
|
return await self.execute_query(sql, params)
|
|
|
|
|
|
async def get_user(self, user_id: str) -> Optional[Dict[str, Any]]:
|
|
|
"""Get user by ID"""
|
|
|
|
|
|
sql = "SELECT * FROM users WHERE id = ?"
|
|
|
params = [user_id]
|
|
|
|
|
|
result = await self.execute_query(sql, params)
|
|
|
|
|
|
|
|
|
if result.get("success") and result.get("result"):
|
|
|
rows = result["result"][0].get("results", [])
|
|
|
if rows:
|
|
|
user = rows[0]
|
|
|
if user.get("metadata"):
|
|
|
import json
|
|
|
|
|
|
user["metadata"] = json.loads(user["metadata"])
|
|
|
return user
|
|
|
|
|
|
return None
|
|
|
|
|
|
async def get_user_by_username(self, username: str) -> Optional[Dict[str, Any]]:
|
|
|
"""Get user by username"""
|
|
|
|
|
|
sql = "SELECT * FROM users WHERE username = ?"
|
|
|
params = [username]
|
|
|
|
|
|
result = await self.execute_query(sql, params)
|
|
|
|
|
|
if result.get("success") and result.get("result"):
|
|
|
rows = result["result"][0].get("results", [])
|
|
|
if rows:
|
|
|
user = rows[0]
|
|
|
if user.get("metadata"):
|
|
|
import json
|
|
|
|
|
|
user["metadata"] = json.loads(user["metadata"])
|
|
|
return user
|
|
|
|
|
|
return None
|
|
|
|
|
|
|
|
|
async def create_session(
|
|
|
self,
|
|
|
session_id: str,
|
|
|
user_id: str,
|
|
|
session_data: Dict[str, Any],
|
|
|
expires_at: Optional[int] = None,
|
|
|
) -> Dict[str, Any]:
|
|
|
"""Create a new session"""
|
|
|
|
|
|
sql = """
|
|
|
INSERT INTO sessions (id, user_id, session_data, expires_at)
|
|
|
VALUES (?, ?, ?, ?)
|
|
|
"""
|
|
|
|
|
|
import json
|
|
|
|
|
|
params = [session_id, user_id, json.dumps(session_data), expires_at]
|
|
|
|
|
|
return await self.execute_query(sql, params)
|
|
|
|
|
|
async def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
|
|
|
"""Get session by ID"""
|
|
|
|
|
|
sql = """
|
|
|
SELECT * FROM sessions
|
|
|
WHERE id = ? AND (expires_at IS NULL OR expires_at > strftime('%s', 'now'))
|
|
|
"""
|
|
|
params = [session_id]
|
|
|
|
|
|
result = await self.execute_query(sql, params)
|
|
|
|
|
|
if result.get("success") and result.get("result"):
|
|
|
rows = result["result"][0].get("results", [])
|
|
|
if rows:
|
|
|
session = rows[0]
|
|
|
if session.get("session_data"):
|
|
|
import json
|
|
|
|
|
|
session["session_data"] = json.loads(session["session_data"])
|
|
|
return session
|
|
|
|
|
|
return None
|
|
|
|
|
|
async def delete_session(self, session_id: str) -> Dict[str, Any]:
|
|
|
"""Delete a session"""
|
|
|
|
|
|
sql = "DELETE FROM sessions WHERE id = ?"
|
|
|
params = [session_id]
|
|
|
|
|
|
return await self.execute_query(sql, params)
|
|
|
|
|
|
|
|
|
async def create_conversation(
|
|
|
self,
|
|
|
conversation_id: str,
|
|
|
user_id: str,
|
|
|
title: Optional[str] = None,
|
|
|
messages: Optional[List[Dict[str, Any]]] = None,
|
|
|
) -> Dict[str, Any]:
|
|
|
"""Create a new conversation"""
|
|
|
|
|
|
sql = """
|
|
|
INSERT INTO conversations (id, user_id, title, messages)
|
|
|
VALUES (?, ?, ?, ?)
|
|
|
"""
|
|
|
|
|
|
import json
|
|
|
|
|
|
params = [conversation_id, user_id, title, json.dumps(messages or [])]
|
|
|
|
|
|
return await self.execute_query(sql, params)
|
|
|
|
|
|
async def get_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]:
|
|
|
"""Get conversation by ID"""
|
|
|
|
|
|
sql = "SELECT * FROM conversations WHERE id = ?"
|
|
|
params = [conversation_id]
|
|
|
|
|
|
result = await self.execute_query(sql, params)
|
|
|
|
|
|
if result.get("success") and result.get("result"):
|
|
|
rows = result["result"][0].get("results", [])
|
|
|
if rows:
|
|
|
conversation = rows[0]
|
|
|
if conversation.get("messages"):
|
|
|
import json
|
|
|
|
|
|
conversation["messages"] = json.loads(conversation["messages"])
|
|
|
return conversation
|
|
|
|
|
|
return None
|
|
|
|
|
|
async def update_conversation_messages(
|
|
|
self, conversation_id: str, messages: List[Dict[str, Any]]
|
|
|
) -> Dict[str, Any]:
|
|
|
"""Update conversation messages"""
|
|
|
|
|
|
sql = """
|
|
|
UPDATE conversations
|
|
|
SET messages = ?, updated_at = strftime('%s', 'now')
|
|
|
WHERE id = ?
|
|
|
"""
|
|
|
|
|
|
import json
|
|
|
|
|
|
params = [json.dumps(messages), conversation_id]
|
|
|
|
|
|
return await self.execute_query(sql, params)
|
|
|
|
|
|
async def get_user_conversations(
|
|
|
self, user_id: str, limit: int = 50
|
|
|
) -> List[Dict[str, Any]]:
|
|
|
"""Get user's conversations"""
|
|
|
|
|
|
sql = """
|
|
|
SELECT id, user_id, title, created_at, updated_at
|
|
|
FROM conversations
|
|
|
WHERE user_id = ?
|
|
|
ORDER BY updated_at DESC
|
|
|
LIMIT ?
|
|
|
"""
|
|
|
params = [user_id, limit]
|
|
|
|
|
|
result = await self.execute_query(sql, params)
|
|
|
|
|
|
if result.get("success") and result.get("result"):
|
|
|
return result["result"][0].get("results", [])
|
|
|
|
|
|
return []
|
|
|
|
|
|
|
|
|
async def create_agent_execution(
|
|
|
self,
|
|
|
execution_id: str,
|
|
|
user_id: str,
|
|
|
session_id: Optional[str] = None,
|
|
|
task_description: Optional[str] = None,
|
|
|
status: str = "pending",
|
|
|
) -> Dict[str, Any]:
|
|
|
"""Create a new agent execution record"""
|
|
|
|
|
|
sql = """
|
|
|
INSERT INTO agent_executions (id, user_id, session_id, task_description, status)
|
|
|
VALUES (?, ?, ?, ?, ?)
|
|
|
"""
|
|
|
|
|
|
params = [execution_id, user_id, session_id, task_description, status]
|
|
|
|
|
|
return await self.execute_query(sql, params)
|
|
|
|
|
|
async def update_agent_execution(
|
|
|
self,
|
|
|
execution_id: str,
|
|
|
status: Optional[str] = None,
|
|
|
result: Optional[str] = None,
|
|
|
execution_time: Optional[int] = None,
|
|
|
) -> Dict[str, Any]:
|
|
|
"""Update agent execution record"""
|
|
|
|
|
|
updates = []
|
|
|
params = []
|
|
|
|
|
|
if status:
|
|
|
updates.append("status = ?")
|
|
|
params.append(status)
|
|
|
|
|
|
if result:
|
|
|
updates.append("result = ?")
|
|
|
params.append(result)
|
|
|
|
|
|
if execution_time is not None:
|
|
|
updates.append("execution_time = ?")
|
|
|
params.append(execution_time)
|
|
|
|
|
|
if status in ["completed", "failed"]:
|
|
|
updates.append("completed_at = strftime('%s', 'now')")
|
|
|
|
|
|
if not updates:
|
|
|
return {"success": True, "message": "No updates provided"}
|
|
|
|
|
|
sql = f"""
|
|
|
UPDATE agent_executions
|
|
|
SET {', '.join(updates)}
|
|
|
WHERE id = ?
|
|
|
"""
|
|
|
params.append(execution_id)
|
|
|
|
|
|
return await self.execute_query(sql, params)
|
|
|
|
|
|
async def get_agent_execution(self, execution_id: str) -> Optional[Dict[str, Any]]:
|
|
|
"""Get agent execution by ID"""
|
|
|
|
|
|
sql = "SELECT * FROM agent_executions WHERE id = ?"
|
|
|
params = [execution_id]
|
|
|
|
|
|
result = await self.execute_query(sql, params)
|
|
|
|
|
|
if result.get("success") and result.get("result"):
|
|
|
rows = result["result"][0].get("results", [])
|
|
|
if rows:
|
|
|
return rows[0]
|
|
|
|
|
|
return None
|
|
|
|
|
|
async def get_user_executions(
|
|
|
self, user_id: str, limit: int = 50
|
|
|
) -> List[Dict[str, Any]]:
|
|
|
"""Get user's agent executions"""
|
|
|
|
|
|
sql = """
|
|
|
SELECT * FROM agent_executions
|
|
|
WHERE user_id = ?
|
|
|
ORDER BY created_at DESC
|
|
|
LIMIT ?
|
|
|
"""
|
|
|
params = [user_id, limit]
|
|
|
|
|
|
result = await self.execute_query(sql, params)
|
|
|
|
|
|
if result.get("success") and result.get("result"):
|
|
|
return result["result"][0].get("results", [])
|
|
|
|
|
|
return []
|
|
|
|
|
|
|
|
|
async def create_file_record(
|
|
|
self,
|
|
|
file_id: str,
|
|
|
user_id: str,
|
|
|
filename: str,
|
|
|
file_key: str,
|
|
|
file_size: int,
|
|
|
content_type: str,
|
|
|
bucket: str = "storage",
|
|
|
) -> Dict[str, Any]:
|
|
|
"""Create a file record"""
|
|
|
|
|
|
sql = """
|
|
|
INSERT INTO files (id, user_id, filename, file_key, file_size, content_type, bucket)
|
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
|
"""
|
|
|
|
|
|
params = [file_id, user_id, filename, file_key, file_size, content_type, bucket]
|
|
|
|
|
|
return await self.execute_query(sql, params)
|
|
|
|
|
|
async def get_file_record(self, file_id: str) -> Optional[Dict[str, Any]]:
|
|
|
"""Get file record by ID"""
|
|
|
|
|
|
sql = "SELECT * FROM files WHERE id = ?"
|
|
|
params = [file_id]
|
|
|
|
|
|
result = await self.execute_query(sql, params)
|
|
|
|
|
|
if result.get("success") and result.get("result"):
|
|
|
rows = result["result"][0].get("results", [])
|
|
|
if rows:
|
|
|
return rows[0]
|
|
|
|
|
|
return None
|
|
|
|
|
|
async def get_user_files(
|
|
|
self, user_id: str, limit: int = 100
|
|
|
) -> List[Dict[str, Any]]:
|
|
|
"""Get user's files"""
|
|
|
|
|
|
sql = """
|
|
|
SELECT * FROM files
|
|
|
WHERE user_id = ?
|
|
|
ORDER BY created_at DESC
|
|
|
LIMIT ?
|
|
|
"""
|
|
|
params = [user_id, limit]
|
|
|
|
|
|
result = await self.execute_query(sql, params)
|
|
|
|
|
|
if result.get("success") and result.get("result"):
|
|
|
return result["result"][0].get("results", [])
|
|
|
|
|
|
return []
|
|
|
|
|
|
async def delete_file_record(self, file_id: str) -> Dict[str, Any]:
|
|
|
"""Delete a file record"""
|
|
|
|
|
|
sql = "DELETE FROM files WHERE id = ?"
|
|
|
params = [file_id]
|
|
|
|
|
|
return await self.execute_query(sql, params)
|
|
|
|
|
|
|
|
|
async def initialize_schema(self) -> Dict[str, Any]:
|
|
|
"""Initialize database schema"""
|
|
|
|
|
|
schema_queries = [
|
|
|
{
|
|
|
"sql": """CREATE TABLE IF NOT EXISTS users (
|
|
|
id TEXT PRIMARY KEY,
|
|
|
username TEXT UNIQUE NOT NULL,
|
|
|
email TEXT UNIQUE,
|
|
|
created_at INTEGER DEFAULT (strftime('%s', 'now')),
|
|
|
updated_at INTEGER DEFAULT (strftime('%s', 'now')),
|
|
|
metadata TEXT
|
|
|
)"""
|
|
|
},
|
|
|
{
|
|
|
"sql": """CREATE TABLE IF NOT EXISTS sessions (
|
|
|
id TEXT PRIMARY KEY,
|
|
|
user_id TEXT NOT NULL,
|
|
|
session_data TEXT,
|
|
|
created_at INTEGER DEFAULT (strftime('%s', 'now')),
|
|
|
expires_at INTEGER,
|
|
|
FOREIGN KEY (user_id) REFERENCES users(id)
|
|
|
)"""
|
|
|
},
|
|
|
{
|
|
|
"sql": """CREATE TABLE IF NOT EXISTS conversations (
|
|
|
id TEXT PRIMARY KEY,
|
|
|
user_id TEXT NOT NULL,
|
|
|
title TEXT,
|
|
|
messages TEXT,
|
|
|
created_at INTEGER DEFAULT (strftime('%s', 'now')),
|
|
|
updated_at INTEGER DEFAULT (strftime('%s', 'now')),
|
|
|
FOREIGN KEY (user_id) REFERENCES users(id)
|
|
|
)"""
|
|
|
},
|
|
|
{
|
|
|
"sql": """CREATE TABLE IF NOT EXISTS files (
|
|
|
id TEXT PRIMARY KEY,
|
|
|
user_id TEXT NOT NULL,
|
|
|
filename TEXT NOT NULL,
|
|
|
file_key TEXT NOT NULL,
|
|
|
file_size INTEGER,
|
|
|
content_type TEXT,
|
|
|
bucket TEXT DEFAULT 'storage',
|
|
|
created_at INTEGER DEFAULT (strftime('%s', 'now')),
|
|
|
FOREIGN KEY (user_id) REFERENCES users(id)
|
|
|
)"""
|
|
|
},
|
|
|
{
|
|
|
"sql": """CREATE TABLE IF NOT EXISTS agent_executions (
|
|
|
id TEXT PRIMARY KEY,
|
|
|
user_id TEXT NOT NULL,
|
|
|
session_id TEXT,
|
|
|
task_description TEXT,
|
|
|
status TEXT DEFAULT 'pending',
|
|
|
result TEXT,
|
|
|
execution_time INTEGER,
|
|
|
created_at INTEGER DEFAULT (strftime('%s', 'now')),
|
|
|
completed_at INTEGER,
|
|
|
FOREIGN KEY (user_id) REFERENCES users(id)
|
|
|
)"""
|
|
|
},
|
|
|
]
|
|
|
|
|
|
|
|
|
index_queries = [
|
|
|
{
|
|
|
"sql": "CREATE INDEX IF NOT EXISTS idx_sessions_user_id ON sessions(user_id)"
|
|
|
},
|
|
|
{
|
|
|
"sql": "CREATE INDEX IF NOT EXISTS idx_conversations_user_id ON conversations(user_id)"
|
|
|
},
|
|
|
{"sql": "CREATE INDEX IF NOT EXISTS idx_files_user_id ON files(user_id)"},
|
|
|
{
|
|
|
"sql": "CREATE INDEX IF NOT EXISTS idx_agent_executions_user_id ON agent_executions(user_id)"
|
|
|
},
|
|
|
]
|
|
|
|
|
|
all_queries = schema_queries + index_queries
|
|
|
|
|
|
return await self.batch_execute(all_queries)
|
|
|
|
