MrKalzu/sliding_sqlite
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

First push to rauhala gitea.

Browse Source
This commit is contained in:
Kalzu Rekku
2025-03-16 10:30:47 +02:00
parent 2a3613280f
commit dc5d652503
6 changed files with 1597 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
SlidingSqlite.py
+720
View File
@@ -0,0 +1,720 @@
import sqlite3
import uuid
import threading
import datetime
import queue
import os
import logging
import time
from typing import (
Any,
Dict,
Optional,
Tuple,
List,
Union,
Set,
NamedTuple,
TypeVar,
Generic,
Callable,
)
T = TypeVar("T")
class DatabaseError(Exception):
"""Base exception for all database errors"""
pass
class QueryError(DatabaseError):
"""Exception raised when a query fails"""
pass
class QueryResult(Generic[T]):
"""Class to handle query results with proper error handling"""
def __init__(self, data: Optional[T] = None, error: Optional[Exception] = None):
self.data = data
self.error = error
self.success = error is None and data is not None
def __bool__(self) -> bool:
return self.success
class DatabaseTimeframe(NamedTuple):
"""Represents a database file and its time range"""
db_file: str
start_time: float
end_time: float
class QueryLockManager:
def __init__(self, sliding_sqlite, query_id):
self.sliding_sqlite = sliding_sqlite
self.query_id = query_id
self.lock = sliding_sqlite.query_lock
self.is_active = False
def __enter__(self):
# Acquire the lock and check query status
with self.lock:
self.is_active = (
self.query_id in self.sliding_sqlite.read_queues
and self.query_id in self.sliding_sqlite.active_queries
)
return self.is_active
def __exit__(self, exc_type, exc_val, exc_tb):
# If there was an exception, we don't need to do anything
pass
class SlidingSQLite:
"""
Thread-safe SQLite implementation with automatic time-based database rotation.
This class provides a way to safely use SQLite in a multi-threaded environment
by queuing database operations and processing them in dedicated worker threads.
Databases are created based on the specified rotation interval and old databases
are automatically cleaned up based on the specified retention period.
"""
def __init__(
self,
db_dir: str,
schema: str,
retention_period: int = 604800,
rotation_interval: int = 3600,
cleanup_interval: int = 3600,
auto_delete_old_dbs: bool = True,
) -> None:
"""
Initialize the SlidingSQLite instance.
Args:
db_dir: Directory to store database files
schema: SQL schema to initialize new databases
retention_period: Number of seconds to keep databases before deletion (default: 7 days)
rotation_interval: How often to rotate to a new database file in seconds (default: 1 hour)
cleanup_interval: How often to run the cleanup process in seconds (default: 1 hour)
auto_delete_old_dbs: Whether to automatically delete old databases (default: True)
"""
self.db_dir = db_dir
self.schema = schema
self.retention_period = retention_period # In seconds
self.rotation_interval = rotation_interval # In seconds
self.cleanup_interval = cleanup_interval # In seconds
self.auto_delete_old_dbs = auto_delete_old_dbs # New field
# Queues for operations
self.write_queue: queue.Queue[Tuple[str, Tuple[Any, ...], uuid.UUID]] = (
queue.Queue()
)
self.result_queues: Dict[uuid.UUID, queue.Queue[QueryResult[bool]]] = {}
self.read_queues: Dict[
uuid.UUID, queue.Queue[QueryResult[List[Tuple[Any, ...]]]]
] = {}
# Thread synchronization
self.shutdown_flag = threading.Event()
self.worker_thread = None
# Cache for database connections
self.connections: Dict[str, sqlite3.Connection] = {}
self.conn_lock = threading.Lock()
# Track active query IDs for cleanup
self.active_queries: Set[uuid.UUID] = set()
self.query_lock = threading.Lock()
# Initialize system
self._setup()
def _setup(self) -> None:
"""Setup the database directory and initialize workers"""
try:
os.makedirs(self.db_dir, exist_ok=True)
self._init_metadata()
# Start worker threads
self._start_worker()
self._start_cleanup_worker()
# Register current database
self._register_current_db()
except Exception as e:
logging.error(f"Failed to initialize SlidingSQLite: {e}")
raise DatabaseError(f"Failed to initialize SlidingSQLite: {e}")
def _init_metadata(self) -> None:
"""Initialize the metadata database"""
try:
with self._get_connection(self._get_metadata_db()) as conn:
conn.execute(
"""
CREATE TABLE IF NOT EXISTS metadata (
id INTEGER PRIMARY KEY AUTOINCREMENT,
db_file TEXT NOT NULL UNIQUE,
start_time REAL NOT NULL,
end_time REAL NOT NULL
)
"""
)
conn.commit()
except sqlite3.Error as e:
logging.error(f"Failed to initialize metadata database: {e}")
raise DatabaseError(f"Failed to initialize metadata database: {e}")
def _get_connection(self, db_file: str) -> sqlite3.Connection:
"""
Get a connection to the specified database file.
Reuses existing connections when possible.
Args:
db_file: Path to the database file
Returns:
SQLite connection object
"""
with self.conn_lock:
if db_file not in self.connections or self.connections[db_file] is None:
try:
conn = sqlite3.connect(
db_file, isolation_level=None, check_same_thread=False
)
conn.execute("PRAGMA journal_mode=WAL;")
conn.execute(
"PRAGMA busy_timeout=5000;"
) # Wait up to 5 seconds when database is locked
# Initialize schema if this is a data database (not metadata)
if db_file != self._get_metadata_db():
conn.executescript(self.schema)
self.connections[db_file] = conn
except sqlite3.Error as e:
logging.error(f"Failed to connect to database {db_file}: {e}")
raise DatabaseError(f"Failed to connect to database {db_file}: {e}")
return self.connections[db_file]
def _get_metadata_db(self) -> str:
"""Get the path to the metadata database"""
return os.path.join(self.db_dir, "metadata.db")
def _get_current_db(self) -> str:
"""Get the path to the current time-based database"""
# Generate timestamped DB name based on rotation interval
now = time.time()
interval_timestamp = int(now // self.rotation_interval) * self.rotation_interval
timestamp_str = datetime.datetime.fromtimestamp(interval_timestamp).strftime(
"%Y%m%d_%H%M%S"
)
return os.path.join(self.db_dir, f"data_{timestamp_str}.db")
def _register_current_db(self) -> None:
"""Register the current database in the metadata table"""
current_db = self._get_current_db()
now = time.time()
# Calculate time boundaries for the current database
interval_start = int(now // self.rotation_interval) * self.rotation_interval
interval_end = interval_start + self.rotation_interval
try:
with self._get_connection(self._get_metadata_db()) as conn:
# Check if this database is already registered
existing = conn.execute(
"SELECT id FROM metadata WHERE db_file = ?", (current_db,)
).fetchone()
if not existing:
conn.execute(
"INSERT INTO metadata (db_file, start_time, end_time) VALUES (?, ?, ?)",
(current_db, interval_start, interval_end),
)
conn.commit()
except sqlite3.Error as e:
logging.error(f"Failed to register current database: {e}")
# Continue execution as this is not critical
def _rotate_databases(self) -> None:
"""Delete databases that are older than the retention period"""
if not self.auto_delete_old_dbs:
return # Skip deletion if auto-delete is disabled
cutoff_time = time.time() - self.retention_period
try:
with self._get_connection(self._get_metadata_db()) as conn:
# Find databases older than the cutoff time
old_dbs = conn.execute(
"SELECT db_file FROM metadata WHERE end_time < ?", (cutoff_time,)
).fetchall()
# Delete old database files
for (db_file,) in old_dbs:
self._delete_database_file(db_file)
# Clean up metadata entries
conn.execute("DELETE FROM metadata WHERE end_time < ?", (cutoff_time,))
conn.commit()
except sqlite3.Error as e:
logging.error(f"Database rotation error: {e}")
def _cleanup_stale_queries(self) -> None:
"""Clean up stale query results to prevent memory leaks"""
with self.query_lock:
# Find completed queries to clean up
completed_queries = set()
for query_id in list(self.result_queues.keys()):
if query_id not in self.active_queries:
completed_queries.add(query_id)
for query_id in list(self.read_queues.keys()):
if query_id not in self.active_queries:
completed_queries.add(query_id)
# Remove completed queries from dictionaries
for query_id in completed_queries:
if query_id in self.result_queues:
del self.result_queues[query_id]
if query_id in self.read_queues:
del self.read_queues[query_id]
def _delete_database_file(self, db_file: str) -> bool:
"""
Delete a database file and clean up resources.
Args:
db_file: Path to the database file to delete
Returns:
True if deleted successfully, False otherwise
"""
# Close and remove connection if it exists
with self.conn_lock:
if db_file in self.connections:
try:
self.connections[db_file].close()
except sqlite3.Error:
pass # Ignore errors on close
del self.connections[db_file]
# Delete the file
if os.path.exists(db_file):
try:
os.remove(db_file)
logging.info(f"Deleted database: {db_file}")
return True
except OSError as e:
logging.error(f"Failed to delete database {db_file}: {e}")
return False
return False # File didn't exist
def set_retention_period(self, seconds: int) -> None:
"""
Set the retention period for databases.
Args:
seconds: Number of seconds to keep databases before deletion
"""
self.retention_period = max(0, seconds) # Ensure positive value
def set_auto_delete(self, enabled: bool) -> None:
"""
Enable or disable automatic deletion of old databases.
Args:
enabled: Whether to automatically delete old databases
"""
self.auto_delete_old_dbs = enabled
def delete_databases_before(self, timestamp: float) -> int:
"""
Delete all databases with end_time before the specified timestamp.
Args:
timestamp: Unix timestamp (seconds since epoch)
Returns:
Number of databases deleted
"""
count = 0
try:
with self._get_connection(self._get_metadata_db()) as conn:
# Find databases older than the specified time
old_dbs = conn.execute(
"SELECT db_file FROM metadata WHERE end_time < ?", (timestamp,)
).fetchall()
# Delete old database files
for (db_file,) in old_dbs:
if self._delete_database_file(db_file):
count += 1
# Clean up metadata entries
conn.execute("DELETE FROM metadata WHERE end_time < ?", (timestamp,))
conn.commit()
except sqlite3.Error as e:
logging.error(f"Database deletion error: {e}")
raise DatabaseError(f"Failed to delete databases: {e}")
return count
def delete_databases_in_range(self, start_time: float, end_time: float) -> int:
"""
Delete all databases with time ranges falling within the specified period.
Args:
start_time: Start of time range (unix timestamp)
end_time: End of time range (unix timestamp)
Returns:
Number of databases deleted
"""
count = 0
try:
with self._get_connection(self._get_metadata_db()) as conn:
# Find databases in the specified time range
# A database is in range if its time range overlaps with the specified range
dbs = conn.execute(
"""
SELECT db_file FROM metadata
WHERE (start_time <= ? AND end_time >= ?) OR
(start_time <= ? AND end_time >= ?) OR
(start_time >= ? AND end_time <= ?)
""",
(end_time, start_time, end_time, start_time, start_time, end_time),
).fetchall()
# Delete database files
for (db_file,) in dbs:
if self._delete_database_file(db_file):
count += 1
# Clean up metadata entries
conn.execute(
"""
DELETE FROM metadata
WHERE (start_time <= ? AND end_time >= ?) OR
(start_time <= ? AND end_time >= ?) OR
(start_time >= ? AND end_time <= ?)
""",
(end_time, start_time, end_time, start_time, start_time, end_time),
)
conn.commit()
except sqlite3.Error as e:
logging.error(f"Database deletion error: {e}")
raise DatabaseError(f"Failed to delete databases: {e}")
return count
def get_databases_info(self) -> List[DatabaseTimeframe]:
"""
Get information about all available databases.
Returns:
List of DatabaseTimeframe objects containing database file paths and time ranges
"""
databases = []
try:
with self._get_connection(self._get_metadata_db()) as conn:
rows = conn.execute(
"SELECT db_file, start_time, end_time FROM metadata ORDER BY start_time"
).fetchall()
for db_file, start_time, end_time in rows:
databases.append(DatabaseTimeframe(db_file, start_time, end_time))
except sqlite3.Error as e:
logging.error(f"Error retrieving database info: {e}")
raise DatabaseError(f"Failed to retrieve database info: {e}")
return databases
def execute(self, query: str, params: Tuple[Any, ...] = ()) -> uuid.UUID:
"""
Smart query executor that automatically determines if the query
is a read or write operation and routes accordingly.
Args:
query: SQL query to execute
params: Parameters for the query
Returns:
UUID that can be used to retrieve the result
"""
# look for new database files
self._register_current_db()
query_upper = query.strip().upper()
# Check if the query is a read operation
if (
query_upper.startswith("SELECT")
or query_upper.startswith("PRAGMA")
or query_upper.startswith("EXPLAIN")
):
return self.execute_read(query, params)
else:
return self.execute_write(query, params)
def execute_write(self, query: str, params: Tuple[Any, ...] = ()) -> uuid.UUID:
"""
Execute a write query asynchronously.
Args:
query: SQL query to execute
params: Parameters for the query
Returns:
UUID that can be used to retrieve the result
"""
# look for new database files
self._register_current_db()
query_id = uuid.uuid4()
with self.query_lock:
self.result_queues[query_id] = queue.Queue()
self.active_queries.add(query_id)
self.write_queue.put((query, params, query_id))
return query_id
def execute_write_sync(
self, query: str, params: Tuple[Any, ...] = (), timeout: float = 5.0
) -> QueryResult[bool]:
"""
Execute a write query synchronously.
Args:
query: SQL query to execute
params: Parameters for the query
timeout: Maximum time to wait for the result
Returns:
QueryResult containing success status and any error
"""
query_id = self.execute_write(query, params)
return self.get_result(query_id, timeout)
def execute_read(self, query: str, params: Tuple[Any, ...] = ()) -> uuid.UUID:
"""
Execute a read query asynchronously across all relevant databases.
This provides transparent access to all time-windowed data.
Args:
query: SQL query to execute
params: Parameters for the query
Returns:
UUID that can be used to retrieve the result
"""
# look for new database files
self._register_current_db()
query_id = uuid.uuid4()
with self.query_lock:
self.read_queues[query_id] = queue.Queue()
self.active_queries.add(query_id)
# Start the worker thread that will query across all databases
threading.Thread(
target=self._read_across_all_worker,
args=(query, params, query_id),
daemon=True,
).start()
return query_id
def _read_worker(
self, query: str, params: Tuple[Any, ...], query_id: uuid.UUID
) -> None:
"""Worker thread for processing read queries"""
db_file = self._get_current_db()
try:
with self._get_connection(db_file) as conn:
results = conn.execute(query, params).fetchall()
if query_id in self.read_queues:
self.read_queues[query_id].put(QueryResult(data=results))
except Exception as e:
error_msg = f"Read error: {e}"
logging.error(error_msg)
if query_id in self.read_queues:
self.read_queues[query_id].put(QueryResult(error=QueryError(error_msg)))
def execute_read_sync(
self, query: str, params: Tuple[Any, ...] = (), timeout: float = 5.0
) -> QueryResult[List[Tuple[Any, ...]]]:
"""
Execute a read query synchronously across all relevant databases.
Args:
query: SQL query to execute
params: Parameters for the query
timeout: Maximum time to wait for the result
Returns:
QueryResult containing combined query results and any error
"""
query_id = self.execute_read(query, params)
return self.get_read_result(query_id, timeout)
def _read_across_all_worker(
self, query: str, params: Tuple[Any, ...], query_id: uuid.UUID
) -> None:
"""Worker thread for processing read queries across all databases"""
try:
# Get all available databases from metadata
with self._get_connection(self._get_metadata_db()) as conn:
db_files = conn.execute(
"SELECT db_file FROM metadata ORDER BY end_time DESC"
).fetchall()
all_results: List[Tuple[Any, ...]] = []
for (db_file,) in db_files:
if os.path.exists(db_file):
try:
with self._get_connection(db_file) as conn:
results = conn.execute(query, params).fetchall()
all_results.extend(results)
except sqlite3.Error as e:
logging.warning(f"Error reading from {db_file}: {e}")
# Continue with other databases
# Use the context manager to safely check query status
with QueryLockManager(self, query_id) as is_active:
if is_active:
self.read_queues[query_id].put(QueryResult(data=all_results))
else:
logging.warning(
f"Query ID {query_id} no longer active when trying to return results"
)
except Exception as e:
error_msg = f"Failed to execute query across all databases: {e}"
logging.error(error_msg)
with QueryLockManager(self, query_id) as is_active:
if is_active:
self.read_queues[query_id].put(
QueryResult(error=QueryError(error_msg))
)
def get_result(
self, query_id: uuid.UUID, timeout: float = 5.0
) -> QueryResult[bool]:
"""
Get the result of a write query.
Args:
query_id: UUID returned by execute_write
timeout: Maximum time to wait for the result
Returns:
QueryResult containing success status and any error
"""
if query_id not in self.result_queues:
return QueryResult(error=QueryError("Invalid query ID"))
try:
result = self.result_queues[query_id].get(timeout=timeout)
with self.query_lock:
if query_id in self.active_queries:
self.active_queries.remove(query_id)
return result
except queue.Empty:
return QueryResult(error=QueryError("Query timed out"))
def get_read_result(
self, query_id: uuid.UUID, timeout: float = 5.0
) -> QueryResult[List[Tuple[Any, ...]]]:
"""
Get the result of a read query.
Args:
query_id: UUID returned by execute_read
timeout: Maximum time to wait for the result
Returns:
QueryResult containing query results and any error
"""
# Check if the query ID exists in read_queues
with self.query_lock:
if query_id not in self.read_queues:
return QueryResult(error=QueryError("Invalid query ID"))
if query_id not in self.active_queries:
self.active_queries.add(query_id) # Re-add if it was removed
try:
result = self.read_queues[query_id].get(timeout=timeout)
with self.query_lock:
if query_id in self.active_queries:
self.active_queries.remove(query_id)
return result
except queue.Empty:
return QueryResult(error=QueryError("Query timed out"))
def _start_worker(self) -> None:
"""Start the background worker thread for processing write operations."""
if self.worker_thread and self.worker_thread.is_alive():
return
def worker() -> None:
while not self.shutdown_flag.is_set():
try:
task = self.write_queue.get(timeout=1) # Adjust timeout as needed
if task:
self._process_write_task(task)
except queue.Empty:
continue
except Exception as e:
logging.error(f"Worker thread encountered an error: {e}")
self.worker_thread = threading.Thread(target=worker, daemon=True)
self.worker_thread.start()
def _start_cleanup_worker(self) -> None:
"""Start the cleanup worker thread for database rotation."""
threading.Thread(target=self._cleanup_worker, daemon=True).start()
def _process_write_task(self, task: Tuple[str, Tuple[Any, ...], uuid.UUID]) -> None:
"""Process a single write task from the queue."""
query, params, query_id = task
db_file = self._get_current_db()
try:
with self._get_connection(db_file) as conn:
conn.execute(query, params)
conn.commit()
self.result_queues[query_id].put(QueryResult(data=True))
except Exception as e:
logging.error(f"Write error: {e}")
self.result_queues[query_id].put(QueryResult(error=e))
def _cleanup_worker(self) -> None:
"""Worker thread for handling database rotation and cleanup."""
while not self.shutdown_flag.is_set():
self._rotate_databases()
self._cleanup_stale_queries() # Also clean up stale queries
time.sleep(self.cleanup_interval) # Use the configurable interval
def shutdown(self) -> None:
"""Gracefully shut down the workers and close connections."""
self.shutdown_flag.set()
if self.worker_thread:
self.worker_thread.join()
logging.info("SlidingSQLite shutdown completed.")