fix: sqlite conn

Summary:

Test Plan:
This commit is contained in:
Eric Huang 2025-02-26 14:08:17 -08:00
parent c8a20b8ed0
commit 80fba6247d

View file

@ -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."""