Coverage for src / osiris_cli / database_manager.py: 0%
215 statements
« prev ^ index » next coverage.py v7.13.0, created at 2025-12-27 17:41 +0200
1"""
2Database Integration for Osiris CLI
4Unified interface for multiple database types with AI-friendly tools.
5Supports PostgreSQL, MySQL, SQLite, and MongoDB.
6"""
8import json
9from typing import Optional, Dict, List, Any, Union
10from pathlib import Path
11from dataclasses import dataclass
12from enum import Enum
14from .logger import get_logger
15from .errors import DatabaseError, DatabaseConnectionError, DatabaseQueryError
17logger = get_logger()
20class DatabaseType(Enum):
21 """Supported database types"""
22 POSTGRESQL = "postgresql"
23 MYSQL = "mysql"
24 SQLITE = "sqlite"
25 MONGODB = "mongodb"
28@dataclass
29class DatabaseConnection:
30 """Database connection configuration"""
31 id: str
32 type: DatabaseType
33 host: Optional[str] = None
34 port: Optional[int] = None
35 database: Optional[str] = None
36 username: Optional[str] = None
37 password: Optional[str] = None
38 connection_string: Optional[str] = None
40 # SQLite-specific
41 file_path: Optional[Path] = None
43 # MongoDB-specific
44 auth_source: Optional[str] = None
46 def __post_init__(self):
47 """Validate configuration"""
48 if self.type == DatabaseType.SQLITE:
49 if not self.file_path:
50 raise ValueError("SQLite requires file_path")
51 elif self.type == DatabaseType.MONGODB:
52 if not self.connection_string and not (self.host and self.database):
53 raise ValueError("MongoDB requires connection_string or host+database")
54 else: # PostgreSQL, MySQL
55 if not self.connection_string and not (self.host and self.database):
56 raise ValueError(f"{self.type.value} requires connection_string or host+database")
59class DatabaseManager:
60 """
61 Unified database manager for multiple database types.
63 Features:
64 - Connection management
65 - Schema exploration
66 - Safe query execution
67 - Result formatting
68 - Transaction support
69 """
71 def __init__(self):
72 """Initialize database manager"""
73 self.connections: Dict[str, Any] = {}
74 self.configs: Dict[str, DatabaseConnection] = {}
75 logger.info("Database manager initialized")
77 def connect(
78 self,
79 connection_id: str,
80 db_type: str,
81 **kwargs
82 ) -> bool:
83 """
84 Connect to a database.
86 Args:
87 connection_id: Unique connection identifier
88 db_type: Database type (postgresql, mysql, sqlite, mongodb)
89 **kwargs: Connection parameters
91 Returns:
92 True if connected successfully
93 """
94 try:
95 db_type_enum = DatabaseType(db_type.lower())
97 # Create connection config
98 config = DatabaseConnection(
99 id=connection_id,
100 type=db_type_enum,
101 **kwargs
102 )
104 # Establish connection based on type
105 if db_type_enum == DatabaseType.POSTGRESQL:
106 connection = self._connect_postgresql(config)
107 elif db_type_enum == DatabaseType.MYSQL:
108 connection = self._connect_mysql(config)
109 elif db_type_enum == DatabaseType.SQLITE:
110 connection = self._connect_sqlite(config)
111 elif db_type_enum == DatabaseType.MONGODB:
112 connection = self._connect_mongodb(config)
113 else:
114 raise ValueError(f"Unsupported database type: {db_type}")
116 self.connections[connection_id] = connection
117 self.configs[connection_id] = config
119 logger.info(f"Connected to {db_type} database: {connection_id}")
120 return True
122 except Exception as e:
123 logger.error(f"Failed to connect to {db_type}: {e}", exc_info=True)
124 raise DatabaseConnectionError(db_type, str(e))
126 def _connect_postgresql(self, config: DatabaseConnection) -> Any:
127 """Connect to PostgreSQL"""
128 try:
129 import psycopg2
131 if config.connection_string:
132 conn = psycopg2.connect(config.connection_string)
133 else:
134 conn = psycopg2.connect(
135 host=config.host,
136 port=config.port or 5432,
137 database=config.database,
138 user=config.username,
139 password=config.password
140 )
142 return conn
144 except ImportError:
145 raise DatabaseConnectionError(
146 "PostgreSQL",
147 "psycopg2 not installed. Install with: pip install psycopg2-binary"
148 )
150 def _connect_mysql(self, config: DatabaseConnection) -> Any:
151 """Connect to MySQL"""
152 try:
153 import mysql.connector
155 if config.connection_string:
156 # Parse connection string
157 conn = mysql.connector.connect(config.connection_string)
158 else:
159 conn = mysql.connector.connect(
160 host=config.host,
161 port=config.port or 3306,
162 database=config.database,
163 user=config.username,
164 password=config.password
165 )
167 return conn
169 except ImportError:
170 raise DatabaseConnectionError(
171 "MySQL",
172 "mysql-connector-python not installed. Install with: pip install mysql-connector-python"
173 )
175 def _connect_sqlite(self, config: DatabaseConnection) -> Any:
176 """Connect to SQLite"""
177 import sqlite3
179 file_path = Path(config.file_path)
180 file_path.parent.mkdir(parents=True, exist_ok=True)
182 conn = sqlite3.connect(str(file_path))
183 conn.row_factory = sqlite3.Row # Enable column access by name
185 return conn
187 def _connect_mongodb(self, config: DatabaseConnection) -> Any:
188 """Connect to MongoDB"""
189 try:
190 from pymongo import MongoClient
192 if config.connection_string:
193 client = MongoClient(config.connection_string)
194 else:
195 client = MongoClient(
196 host=config.host,
197 port=config.port or 27017,
198 username=config.username,
199 password=config.password,
200 authSource=config.auth_source or "admin"
201 )
203 # Get database
204 db = client[config.database]
205 return db
207 except ImportError:
208 raise DatabaseConnectionError(
209 "MongoDB",
210 "pymongo not installed. Install with: pip install pymongo"
211 )
213 def disconnect(self, connection_id: str) -> bool:
214 """Disconnect from database"""
215 if connection_id not in self.connections:
216 return False
218 try:
219 conn = self.connections[connection_id]
220 config = self.configs[connection_id]
222 if config.type == DatabaseType.MONGODB:
223 conn.client.close()
224 else:
225 conn.close()
227 del self.connections[connection_id]
228 del self.configs[connection_id]
230 logger.info(f"Disconnected: {connection_id}")
231 return True
233 except Exception as e:
234 logger.error(f"Error disconnecting {connection_id}: {e}")
235 return False
237 def execute_query(
238 self,
239 connection_id: str,
240 query: str,
241 params: Optional[Union[tuple, dict]] = None,
242 safe_mode: bool = True
243 ) -> Dict[str, Any]:
244 """
245 Execute SQL query.
247 Args:
248 connection_id: Connection identifier
249 query: SQL query to execute
250 params: Query parameters
251 safe_mode: If True, only allow SELECT queries
253 Returns:
254 Query result dict
255 """
256 if connection_id not in self.connections:
257 raise DatabaseError(f"Connection not found: {connection_id}")
259 conn = self.connections[connection_id]
260 config = self.configs[connection_id]
262 # Safety check
263 if safe_mode:
264 query_upper = query.strip().upper()
265 if not query_upper.startswith('SELECT') and not query_upper.startswith('EXPLAIN'):
266 raise DatabaseQueryError(
267 query,
268 "Only SELECT and EXPLAIN queries allowed in safe mode"
269 )
271 try:
272 if config.type == DatabaseType.MONGODB:
273 return self._execute_mongodb_query(conn, query)
274 else:
275 return self._execute_sql_query(conn, query, params)
277 except Exception as e:
278 logger.error(f"Query failed: {e}", exc_info=True)
279 raise DatabaseQueryError(query, str(e))
281 def _execute_sql_query(
282 self,
283 conn: Any,
284 query: str,
285 params: Optional[Union[tuple, dict]] = None
286 ) -> Dict[str, Any]:
287 """Execute SQL query (PostgreSQL, MySQL, SQLite)"""
288 cursor = conn.cursor()
290 try:
291 if params:
292 cursor.execute(query, params)
293 else:
294 cursor.execute(query)
296 # Fetch results for SELECT
297 if query.strip().upper().startswith('SELECT'):
298 columns = [desc[0] for desc in cursor.description]
299 rows = cursor.fetchall()
301 # Convert rows to dicts
302 results = [
303 dict(zip(columns, row))
304 for row in rows
305 ]
307 return {
308 "success": True,
309 "row_count": len(results),
310 "columns": columns,
311 "rows": results
312 }
313 else:
314 # For INSERT, UPDATE, DELETE
315 conn.commit()
316 return {
317 "success": True,
318 "rows_affected": cursor.rowcount
319 }
321 finally:
322 cursor.close()
324 def _execute_mongodb_query(self, db: Any, query: str) -> Dict[str, Any]:
325 """Execute MongoDB query (JSON format)"""
326 # Parse JSON query
327 query_doc = json.loads(query)
329 collection_name = query_doc.get("collection")
330 operation = query_doc.get("operation", "find")
331 filter_doc = query_doc.get("filter", {})
333 if not collection_name:
334 raise ValueError("MongoDB query must specify 'collection'")
336 collection = db[collection_name]
338 # Execute operation
339 if operation == "find":
340 limit = query_doc.get("limit", 100)
341 results = list(collection.find(filter_doc).limit(limit))
343 # Convert ObjectId to string
344 for doc in results:
345 if "_id" in doc:
346 doc["_id"] = str(doc["_id"])
348 return {
349 "success": True,
350 "row_count": len(results),
351 "documents": results
352 }
354 elif operation == "count":
355 count = collection.count_documents(filter_doc)
356 return {
357 "success": True,
358 "count": count
359 }
361 else:
362 raise ValueError(f"Unsupported MongoDB operation: {operation}")
364 def get_schema(self, connection_id: str) -> Dict[str, Any]:
365 """Get database schema (tables, columns, types)"""
366 if connection_id not in self.connections:
367 raise DatabaseError(f"Connection not found: {connection_id}")
369 config = self.configs[connection_id]
371 if config.type == DatabaseType.POSTGRESQL:
372 return self._get_postgresql_schema(connection_id)
373 elif config.type == DatabaseType.MYSQL:
374 return self._get_mysql_schema(connection_id)
375 elif config.type == DatabaseType.SQLITE:
376 return self._get_sqlite_schema(connection_id)
377 elif config.type == DatabaseType.MONGODB:
378 return self._get_mongodb_schema(connection_id)
380 def _get_postgresql_schema(self, connection_id: str) -> Dict[str, Any]:
381 """Get PostgreSQL schema"""
382 query = """
383 SELECT
384 table_name,
385 column_name,
386 data_type,
387 is_nullable
388 FROM information_schema.columns
389 WHERE table_schema = 'public'
390 ORDER BY table_name, ordinal_position
391 """
393 result = self.execute_query(connection_id, query, safe_mode=False)
395 # Group by table
396 tables = {}
397 for row in result["rows"]:
398 table = row["table_name"]
399 if table not in tables:
400 tables[table] = []
402 tables[table].append({
403 "column": row["column_name"],
404 "type": row["data_type"],
405 "nullable": row["is_nullable"] == "YES"
406 })
408 return {
409 "database_type": "PostgreSQL",
410 "table_count": len(tables),
411 "tables": tables
412 }
414 def _get_mysql_schema(self, connection_id: str) -> Dict[str, Any]:
415 """Get MySQL schema"""
416 query = """
417 SELECT
418 TABLE_NAME as table_name,
419 COLUMN_NAME as column_name,
420 DATA_TYPE as data_type,
421 IS_NULLABLE as is_nullable
422 FROM information_schema.COLUMNS
423 WHERE TABLE_SCHEMA = DATABASE()
424 ORDER BY TABLE_NAME, ORDINAL_POSITION
425 """
427 result = self.execute_query(connection_id, query, safe_mode=False)
429 # Group by table
430 tables = {}
431 for row in result["rows"]:
432 table = row["table_name"]
433 if table not in tables:
434 tables[table] = []
436 tables[table].append({
437 "column": row["column_name"],
438 "type": row["data_type"],
439 "nullable": row["is_nullable"] == "YES"
440 })
442 return {
443 "database_type": "MySQL",
444 "table_count": len(tables),
445 "tables": tables
446 }
448 def _get_sqlite_schema(self, connection_id: str) -> Dict[str, Any]:
449 """Get SQLite schema"""
450 # Get table list
451 tables_query = "SELECT name FROM sqlite_master WHERE type='table'"
452 tables_result = self.execute_query(connection_id, tables_query, safe_mode=False)
454 tables = {}
455 for row in tables_result["rows"]:
456 table_name = row["name"]
458 # Get columns for this table
459 pragma_query = f"PRAGMA table_info({table_name})"
460 columns_result = self.execute_query(connection_id, pragma_query, safe_mode=False)
462 tables[table_name] = [
463 {
464 "column": col["name"],
465 "type": col["type"],
466 "nullable": col["notnull"] == 0
467 }
468 for col in columns_result["rows"]
469 ]
471 return {
472 "database_type": "SQLite",
473 "table_count": len(tables),
474 "tables": tables
475 }
477 def _get_mongodb_schema(self, connection_id: str) -> Dict[str, Any]:
478 """Get MongoDB schema (collections)"""
479 db = self.connections[connection_id]
481 collections = db.list_collection_names()
483 # Sample documents from each collection to infer schema
484 schema = {}
485 for coll_name in collections:
486 collection = db[coll_name]
488 # Sample first document
489 sample = collection.find_one()
490 if sample:
491 # Get field types
492 fields = {
493 key: type(value).__name__
494 for key, value in sample.items()
495 }
496 schema[coll_name] = fields
498 return {
499 "database_type": "MongoDB",
500 "collection_count": len(collections),
501 "collections": schema
502 }
504 def list_connections(self) -> List[Dict[str, Any]]:
505 """List all active connections"""
506 return [
507 {
508 "id": conn_id,
509 "type": config.type.value,
510 "database": config.database,
511 "host": config.host
512 }
513 for conn_id, config in self.configs.items()
514 ]
517# Global database manager instance
518db_manager = DatabaseManager()