Coverage for src / osiris_cli / database_tools.py: 0%
83 statements
« prev ^ index » next coverage.py v7.13.0, created at 2025-12-27 17:41 +0200
1"""
2Database Tools for AI Integration
4AI-friendly database tools with dashboard UI integration.
5"""
7from typing import Dict, Any, List, Optional
8from pathlib import Path
10from .database_manager import db_manager, DatabaseType
11from .dashboard_ui import dashboard, Icons, BlockState, show_database_operation
12from .logger import get_logger
13from .errors import DatabaseError
15logger = get_logger()
18# Tool definitions for AI
19DATABASE_TOOLS = [
20 {
21 "type": "function",
22 "function": {
23 "name": "db_connect",
24 "description": "Connect to a database (PostgreSQL, MySQL, SQLite, or MongoDB). Returns connection ID for subsequent operations.",
25 "parameters": {
26 "type": "object",
27 "properties": {
28 "connection_id": {
29 "type": "string",
30 "description": "Unique identifier for this connection"
31 },
32 "db_type": {
33 "type": "string",
34 "enum": ["postgresql", "mysql", "sqlite", "mongodb"],
35 "description": "Database type"
36 },
37 "host": {
38 "type": "string",
39 "description": "Database host (not needed for SQLite)"
40 },
41 "port": {
42 "type": "integer",
43 "description": "Database port (optional, uses defaults)"
44 },
45 "database": {
46 "type": "string",
47 "description": "Database name"
48 },
49 "username": {
50 "type": "string",
51 "description": "Database username"
52 },
53 "password": {
54 "type": "string",
55 "description": "Database password"
56 },
57 "file_path": {
58 "type": "string",
59 "description": "File path for SQLite database"
60 }
61 },
62 "required": ["connection_id", "db_type"]
63 }
64 }
65 },
66 {
67 "type": "function",
68 "function": {
69 "name": "db_query",
70 "description": "Execute SQL query on connected database. Safe mode only allows SELECT queries. Returns results as JSON.",
71 "parameters": {
72 "type": "object",
73 "properties": {
74 "connection_id": {
75 "type": "string",
76 "description": "Connection ID from db_connect"
77 },
78 "query": {
79 "type": "string",
80 "description": "SQL query to execute"
81 },
82 "safe_mode": {
83 "type": "boolean",
84 "description": "If true, only SELECT queries allowed (default: true)"
85 }
86 },
87 "required": ["connection_id", "query"]
88 }
89 }
90 },
91 {
92 "type": "function",
93 "function": {
94 "name": "db_schema",
95 "description": "Get database schema (tables, columns, types). Useful for understanding database structure.",
96 "parameters": {
97 "type": "object",
98 "properties": {
99 "connection_id": {
100 "type": "string",
101 "description": "Connection ID from db_connect"
102 }
103 },
104 "required": ["connection_id"]
105 }
106 }
107 },
108 {
109 "type": "function",
110 "function": {
111 "name": "db_disconnect",
112 "description": "Disconnect from database and clean up resources.",
113 "parameters": {
114 "type": "object",
115 "properties": {
116 "connection_id": {
117 "type": "string",
118 "description": "Connection ID to disconnect"
119 }
120 },
121 "required": ["connection_id"]
122 }
123 }
124 },
125 {
126 "type": "function",
127 "function": {
128 "name": "db_list_connections",
129 "description": "List all active database connections.",
130 "parameters": {
131 "type": "object",
132 "properties": {}
133 }
134 }
135 }
136]
139# Tool implementations
140def tool_db_connect(
141 connection_id: str,
142 db_type: str,
143 **kwargs
144) -> Dict[str, Any]:
145 """Connect to database"""
146 # Add checklist item
147 item = dashboard.add_checklist_item(f"Connect to {db_type} database")
149 try:
150 # Attempt connection
151 success = db_manager.connect(connection_id, db_type, **kwargs)
153 if success:
154 item.check()
155 dashboard.update()
157 # Add success block
158 content = f"Database Type: {db_type}\n"
159 if kwargs.get('host'):
160 content += f"Host: {kwargs['host']}\n"
161 if kwargs.get('database'):
162 content += f"Database: {kwargs['database']}\n"
164 dashboard.add_block(
165 title=f"Connected: {connection_id}",
166 content=content,
167 icon=Icons.DATABASE,
168 state=BlockState.COMPLETE
169 )
171 return {
172 "success": True,
173 "connection_id": connection_id,
174 "message": f"Connected to {db_type} database"
175 }
176 else:
177 raise DatabaseError(f"Failed to connect to {db_type}")
179 except Exception as e:
180 logger.error(f"Database connection failed: {e}")
182 # Add error block
183 dashboard.add_block(
184 title=f"Connection Failed: {connection_id}",
185 content=str(e),
186 icon=Icons.ERROR,
187 state=BlockState.ERROR
188 )
190 return {
191 "success": False,
192 "error": str(e)
193 }
196def tool_db_query(
197 connection_id: str,
198 query: str,
199 safe_mode: bool = True
200) -> Dict[str, Any]:
201 """Execute database query"""
202 # Add checklist item
203 item = dashboard.add_checklist_item(f"Execute query on {connection_id}")
205 try:
206 # Execute query
207 result = db_manager.execute_query(connection_id, query, safe_mode=safe_mode)
209 item.check()
210 dashboard.update()
212 # Show in dashboard
213 show_database_operation(
214 operation="Query",
215 database=connection_id,
216 query=query,
217 result=result
218 )
220 return result
222 except Exception as e:
223 logger.error(f"Query failed: {e}")
225 # Add error block
226 dashboard.add_block(
227 title=f"Query Failed: {connection_id}",
228 content=f"Query: {query}\n\nError: {str(e)}",
229 icon=Icons.ERROR,
230 state=BlockState.ERROR
231 )
233 return {
234 "success": False,
235 "error": str(e)
236 }
239def tool_db_schema(connection_id: str) -> Dict[str, Any]:
240 """Get database schema"""
241 # Add checklist item
242 item = dashboard.add_checklist_item(f"Get schema for {connection_id}")
244 try:
245 schema = db_manager.get_schema(connection_id)
247 item.check()
248 dashboard.update()
250 # Format schema for display
251 content = f"Database Type: {schema['database_type']}\n"
252 content += f"Tables: {schema['table_count']}\n\n"
254 # List tables
255 if 'tables' in schema:
256 for table_name, columns in list(schema['tables'].items())[:5]:
257 content += f"\nTable: {table_name}\n"
258 for col in columns[:10]:
259 content += f" - {col['column']} ({col['type']})\n"
261 dashboard.add_block(
262 title=f"Schema: {connection_id}",
263 content=content,
264 icon=Icons.DATABASE,
265 state=BlockState.COLLAPSED
266 )
268 return schema
270 except Exception as e:
271 logger.error(f"Schema fetch failed: {e}")
272 return {
273 "success": False,
274 "error": str(e)
275 }
278def tool_db_disconnect(connection_id: str) -> Dict[str, Any]:
279 """Disconnect from database"""
280 try:
281 success = db_manager.disconnect(connection_id)
283 if success:
284 dashboard.add_message(f"{Icons.SUCCESS} Disconnected from {connection_id}")
285 return {
286 "success": True,
287 "message": f"Disconnected from {connection_id}"
288 }
289 else:
290 return {
291 "success": False,
292 "error": f"Connection {connection_id} not found"
293 }
295 except Exception as e:
296 return {
297 "success": False,
298 "error": str(e)
299 }
302def tool_db_list_connections() -> Dict[str, Any]:
303 """List active connections"""
304 try:
305 connections = db_manager.list_connections()
307 # Format for display
308 content = ""
309 for conn in connections:
310 content += f"{Icons.DATABASE} {conn['id']} - {conn['type']}\n"
311 content += f" Database: {conn['database']}\n"
312 if conn['host']:
313 content += f" Host: {conn['host']}\n"
314 content += "\n"
316 if content:
317 dashboard.add_block(
318 title="Active Database Connections",
319 content=content,
320 icon=Icons.DATABASE,
321 state=BlockState.EXPANDED
322 )
323 else:
324 dashboard.add_message(f"{Icons.INFO} No active connections")
326 return {
327 "success": True,
328 "connections": connections
329 }
331 except Exception as e:
332 return {
333 "success": False,
334 "error": str(e)
335 }
338# Tool registry
339TOOL_HANDLERS = {
340 "db_connect": tool_db_connect,
341 "db_query": tool_db_query,
342 "db_schema": tool_db_schema,
343 "db_disconnect": tool_db_disconnect,
344 "db_list_connections": tool_db_list_connections
345}