40cee6d791
Some checks failed
Cross-Node Transaction Testing / transaction-test (push) Has been cancelled
Deploy to Testnet / deploy-testnet (push) Has been cancelled
Documentation Validation / validate-docs (push) Has been cancelled
Documentation Validation / validate-policies-strict (push) Has been cancelled
Integration Tests / test-service-integration (push) Has been cancelled
Multi-Node Stress Testing / stress-test (push) Has been cancelled
Python Tests / test-python (push) Has been cancelled
Security Scanning / security-scan (push) Has been cancelled
API Endpoint Tests / test-api-endpoints (push) Successful in 20s
CLI Tests / test-cli (push) Failing after 3s
Package Tests / Python package - aitbc-agent-sdk (push) Successful in 33s
Package Tests / Python package - aitbc-core (push) Failing after 1s
Package Tests / Python package - aitbc-crypto (push) Successful in 10s
Package Tests / Python package - aitbc-sdk (push) Successful in 9s
Package Tests / JavaScript package - aitbc-sdk-js (push) Successful in 10s
Package Tests / JavaScript package - aitbc-token (push) Successful in 17s
Production Tests / Production Integration Tests (push) Failing after 6s
- Added List import and field_validator to config.py - Added database connection pooling settings (max_overflow, pool_recycle, pool_pre_ping, echo) - Added rate limiting settings (rate_limit_requests, rate_limit_window_seconds) - Added CORS allow_origins field with default empty list - Added validate_secrets() method to check required secrets in production - Added validate_secret_length() validator for secret_key and jwt_secret (minimum
406 lines
11 KiB
Python
406 lines
11 KiB
Python
"""
|
|
AITBC Database Utilities
|
|
Database connection and query utilities for AITBC applications
|
|
"""
|
|
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
from contextlib import contextmanager
|
|
from .exceptions import DatabaseError
|
|
|
|
# SQLAlchemy support for connection pooling
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
from sqlalchemy.pool import QueuePool, StaticPool
|
|
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
|
|
|
|
|
class DatabaseConnection:
|
|
"""
|
|
Base database connection class for AITBC applications.
|
|
Provides common database operations with error handling.
|
|
"""
|
|
|
|
def __init__(self, db_path: Path, timeout: int = 30):
|
|
"""
|
|
Initialize database connection.
|
|
|
|
Args:
|
|
db_path: Path to database file
|
|
timeout: Connection timeout in seconds
|
|
"""
|
|
self.db_path = db_path
|
|
self.timeout = timeout
|
|
self._connection = None
|
|
|
|
def connect(self) -> sqlite3.Connection:
|
|
"""
|
|
Establish database connection.
|
|
|
|
Returns:
|
|
SQLite connection object
|
|
|
|
Raises:
|
|
DatabaseError: If connection fails
|
|
"""
|
|
try:
|
|
self._connection = sqlite3.connect(
|
|
self.db_path,
|
|
timeout=self.timeout
|
|
)
|
|
self._connection.row_factory = sqlite3.Row
|
|
return self._connection
|
|
except sqlite3.Error as e:
|
|
raise DatabaseError(f"Failed to connect to database: {e}")
|
|
|
|
def close(self) -> None:
|
|
"""Close database connection."""
|
|
if self._connection:
|
|
self._connection.close()
|
|
self._connection = None
|
|
|
|
@contextmanager
|
|
def cursor(self):
|
|
"""
|
|
Context manager for database cursor.
|
|
|
|
Yields:
|
|
Database cursor
|
|
"""
|
|
if not self._connection:
|
|
self.connect()
|
|
cursor = self._connection.cursor()
|
|
try:
|
|
yield cursor
|
|
self._connection.commit()
|
|
except Exception as e:
|
|
self._connection.rollback()
|
|
raise DatabaseError(f"Database operation failed: {e}")
|
|
finally:
|
|
cursor.close()
|
|
|
|
async def execute(
|
|
self,
|
|
query: str,
|
|
params: Optional[Tuple[Any, ...]] = None
|
|
) -> sqlite3.Cursor:
|
|
"""
|
|
Execute a SQL query.
|
|
|
|
Args:
|
|
query: SQL query string
|
|
params: Query parameters
|
|
|
|
Returns:
|
|
Cursor object
|
|
|
|
Raises:
|
|
DatabaseError: If query fails
|
|
"""
|
|
try:
|
|
with self.cursor() as cursor:
|
|
if params:
|
|
cursor.execute(query, params)
|
|
else:
|
|
cursor.execute(query)
|
|
return cursor
|
|
except sqlite3.Error as e:
|
|
raise DatabaseError(f"Query execution failed: {e}")
|
|
|
|
async def fetch_one(
|
|
self,
|
|
query: str,
|
|
params: Optional[Tuple[Any, ...]] = None
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Fetch a single row from query.
|
|
|
|
Args:
|
|
query: SQL query string
|
|
params: Query parameters
|
|
|
|
Returns:
|
|
Row as dictionary or None
|
|
"""
|
|
with self.cursor() as cursor:
|
|
if params:
|
|
cursor.execute(query, params)
|
|
else:
|
|
cursor.execute(query)
|
|
row = cursor.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
async def fetch_all(
|
|
self,
|
|
query: str,
|
|
params: Optional[Tuple[Any, ...]] = None
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Fetch all rows from query.
|
|
|
|
Args:
|
|
query: SQL query string
|
|
params: Query parameters
|
|
|
|
Returns:
|
|
List of rows as dictionaries
|
|
"""
|
|
with self.cursor() as cursor:
|
|
if params:
|
|
cursor.execute(query, params)
|
|
else:
|
|
cursor.execute(query)
|
|
rows = cursor.fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
async def execute_many(
|
|
self,
|
|
query: str,
|
|
params_list: List[Tuple[Any, ...]]
|
|
) -> None:
|
|
"""
|
|
Execute query with multiple parameter sets.
|
|
|
|
Args:
|
|
query: SQL query string
|
|
params_list: List of parameter tuples
|
|
|
|
Raises:
|
|
DatabaseError: If query fails
|
|
"""
|
|
try:
|
|
with self.cursor() as cursor:
|
|
cursor.executemany(query, params_list)
|
|
except sqlite3.Error as e:
|
|
raise DatabaseError(f"Bulk execution failed: {e}")
|
|
|
|
def __enter__(self):
|
|
"""Context manager entry."""
|
|
self.connect()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Context manager exit."""
|
|
self.close()
|
|
|
|
|
|
def get_database_connection(
|
|
db_path: Path,
|
|
timeout: int = 30
|
|
) -> DatabaseConnection:
|
|
"""
|
|
Get a database connection for a given path.
|
|
|
|
Args:
|
|
db_path: Path to database file
|
|
timeout: Connection timeout in seconds
|
|
|
|
Returns:
|
|
DatabaseConnection instance
|
|
"""
|
|
return DatabaseConnection(db_path, timeout)
|
|
|
|
|
|
def ensure_database(db_path: Path) -> Path:
|
|
"""
|
|
Ensure database file and parent directory exist.
|
|
|
|
Args:
|
|
db_path: Path to database file
|
|
|
|
Returns:
|
|
Database path
|
|
"""
|
|
db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
return db_path
|
|
|
|
|
|
def vacuum_database(db_path: Path) -> None:
|
|
"""
|
|
Vacuum database to optimize storage.
|
|
|
|
Args:
|
|
db_path: Path to database file
|
|
|
|
Raises:
|
|
DatabaseError: If vacuum fails
|
|
"""
|
|
try:
|
|
with DatabaseConnection(db_path) as db:
|
|
db.execute("VACUUM")
|
|
except sqlite3.Error as e:
|
|
raise DatabaseError(f"Database vacuum failed: {e}")
|
|
|
|
|
|
def get_table_info(db_path: Path, table_name: str) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get information about a table's columns.
|
|
|
|
Args:
|
|
db_path: Path to database file
|
|
table_name: Name of table
|
|
|
|
Returns:
|
|
List of column information dictionaries
|
|
"""
|
|
with DatabaseConnection(db_path) as db:
|
|
return db.fetch_all(f"PRAGMA table_info({table_name})")
|
|
|
|
|
|
def table_exists(db_path: Path, table_name: str) -> bool:
|
|
"""
|
|
Check if a table exists in the database.
|
|
|
|
Args:
|
|
db_path: Path to database file
|
|
table_name: Name of table
|
|
|
|
Returns:
|
|
True if table exists
|
|
"""
|
|
with DatabaseConnection(db_path) as db:
|
|
result = db.fetch_one(
|
|
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
|
|
(table_name,)
|
|
)
|
|
return result is not None
|
|
|
|
|
|
# SQLAlchemy Connection Pooling Utilities
|
|
|
|
def create_pooled_engine(
|
|
database_url: str,
|
|
pool_size: int = 10,
|
|
max_overflow: int = 20,
|
|
pool_recycle: int = 3600,
|
|
pool_pre_ping: bool = True,
|
|
echo: bool = False,
|
|
use_static_pool: bool = False
|
|
):
|
|
"""
|
|
Create SQLAlchemy engine with connection pooling.
|
|
|
|
Args:
|
|
database_url: Database connection URL
|
|
pool_size: Size of connection pool
|
|
max_overflow: Maximum overflow connections
|
|
pool_recycle: Connection recycle time in seconds
|
|
pool_pre_ping: Test connections before using
|
|
echo: Enable SQL query logging
|
|
use_static_pool: Use StaticPool for SQLite (single connection)
|
|
|
|
Returns:
|
|
SQLAlchemy engine with connection pooling
|
|
"""
|
|
if "sqlite" in database_url and use_static_pool:
|
|
# SQLite with StaticPool (single connection, suitable for tests)
|
|
engine = create_engine(
|
|
database_url,
|
|
connect_args={"check_same_thread": False},
|
|
poolclass=StaticPool,
|
|
echo=echo,
|
|
pool_pre_ping=pool_pre_ping,
|
|
)
|
|
elif "sqlite" in database_url:
|
|
# SQLite with QueuePool (limited pooling support)
|
|
engine = create_engine(
|
|
database_url,
|
|
connect_args={"check_same_thread": False, "timeout": 30},
|
|
poolclass=QueuePool,
|
|
pool_size=min(pool_size, 5), # SQLite has limited concurrent access
|
|
max_overflow=max_overflow,
|
|
pool_pre_ping=pool_pre_ping,
|
|
echo=echo,
|
|
)
|
|
else:
|
|
# PostgreSQL/MySQL with full connection pooling
|
|
engine = create_engine(
|
|
database_url,
|
|
poolclass=QueuePool,
|
|
pool_size=pool_size,
|
|
max_overflow=max_overflow,
|
|
pool_recycle=pool_recycle,
|
|
pool_pre_ping=pool_pre_ping,
|
|
echo=echo,
|
|
)
|
|
return engine
|
|
|
|
|
|
def create_pooled_sessionmaker(
|
|
engine,
|
|
autoflush: bool = False,
|
|
autocommit: bool = False
|
|
):
|
|
"""
|
|
Create session factory with connection pooling.
|
|
|
|
Args:
|
|
engine: SQLAlchemy engine
|
|
autoflush: Enable autoflush
|
|
autocommit: Enable autocommit
|
|
|
|
Returns:
|
|
Session factory
|
|
"""
|
|
return sessionmaker(bind=engine, autoflush=autoflush, autocommit=autocommit)
|
|
|
|
|
|
def create_async_pooled_engine(
|
|
database_url: str,
|
|
pool_size: int = 10,
|
|
max_overflow: int = 20,
|
|
pool_recycle: int = 3600,
|
|
pool_pre_ping: bool = True,
|
|
echo: bool = False
|
|
):
|
|
"""
|
|
Create async SQLAlchemy engine with connection pooling.
|
|
|
|
Args:
|
|
database_url: Database connection URL
|
|
pool_size: Size of connection pool
|
|
max_overflow: Maximum overflow connections
|
|
pool_recycle: Connection recycle time in seconds
|
|
pool_pre_ping: Test connections before using
|
|
echo: Enable SQL query logging
|
|
|
|
Returns:
|
|
Async SQLAlchemy engine with connection pooling
|
|
"""
|
|
# Convert to async URL
|
|
if "sqlite" in database_url:
|
|
async_url = database_url.replace("sqlite:///", "sqlite+aiosqlite:///")
|
|
elif "postgresql" in database_url:
|
|
async_url = database_url.replace("postgresql://", "postgresql+asyncpg://")
|
|
else:
|
|
async_url = database_url
|
|
|
|
engine = create_async_engine(
|
|
async_url,
|
|
poolclass=QueuePool,
|
|
pool_size=pool_size,
|
|
max_overflow=max_overflow,
|
|
pool_recycle=pool_recycle,
|
|
pool_pre_ping=pool_pre_ping,
|
|
echo=echo,
|
|
)
|
|
return engine
|
|
|
|
|
|
def create_async_pooled_sessionmaker(
|
|
engine,
|
|
expire_on_commit: bool = False
|
|
):
|
|
"""
|
|
Create async session factory with connection pooling.
|
|
|
|
Args:
|
|
engine: Async SQLAlchemy engine
|
|
expire_on_commit: Expire objects on commit
|
|
|
|
Returns:
|
|
Async session factory
|
|
"""
|
|
return async_sessionmaker(engine, expire_on_commit=expire_on_commit)
|