mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-08-13 05:17:26 +00:00
fix: sqlite conn
Summary: Test Plan:
This commit is contained in:
parent
c8a20b8ed0
commit
80fba6247d
1 changed files with 19 additions and 9 deletions
|
@ -7,6 +7,7 @@
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import sqlite3
|
import sqlite3
|
||||||
|
import threading
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
||||||
from opentelemetry.sdk.trace import SpanProcessor
|
from opentelemetry.sdk.trace import SpanProcessor
|
||||||
|
@ -17,14 +18,18 @@ class SQLiteSpanProcessor(SpanProcessor):
|
||||||
def __init__(self, conn_string):
|
def __init__(self, conn_string):
|
||||||
"""Initialize the SQLite span processor with a connection string."""
|
"""Initialize the SQLite span processor with a connection string."""
|
||||||
self.conn_string = conn_string
|
self.conn_string = conn_string
|
||||||
self.conn = None
|
self._local = threading.local() # Thread-local storage for connections
|
||||||
self.setup_database()
|
self.setup_database()
|
||||||
|
|
||||||
def _get_connection(self) -> sqlite3.Connection:
|
def _get_connection(self):
|
||||||
"""Get the database connection."""
|
"""Get a thread-local database connection."""
|
||||||
if self.conn is None:
|
if not hasattr(self._local, "conn"):
|
||||||
self.conn = sqlite3.connect(self.conn_string, check_same_thread=False)
|
try:
|
||||||
return self.conn
|
self._local.conn = sqlite3.connect(self.conn_string)
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error connecting to SQLite database: {e}")
|
||||||
|
raise e
|
||||||
|
return self._local.conn
|
||||||
|
|
||||||
def setup_database(self):
|
def setup_database(self):
|
||||||
"""Create the necessary tables if they don't exist."""
|
"""Create the necessary tables if they don't exist."""
|
||||||
|
@ -168,9 +173,14 @@ class SQLiteSpanProcessor(SpanProcessor):
|
||||||
|
|
||||||
def shutdown(self):
|
def shutdown(self):
|
||||||
"""Cleanup any resources."""
|
"""Cleanup any resources."""
|
||||||
if self.conn:
|
# We can't access other threads' connections, so we just close our own
|
||||||
self.conn.close()
|
if hasattr(self._local, "conn"):
|
||||||
self.conn = None
|
try:
|
||||||
|
self._local.conn.close()
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error closing SQLite connection: {e}")
|
||||||
|
finally:
|
||||||
|
del self._local.conn
|
||||||
|
|
||||||
def force_flush(self, timeout_millis=30000):
|
def force_flush(self, timeout_millis=30000):
|
||||||
"""Force export of spans."""
|
"""Force export of spans."""
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue