Spaces:
Build error
Build error
File size: 5,372 Bytes
60e3a80 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import sqlite3
import weakref
from abc import ABC, abstractmethod
from typing import Any, Set
import threading
from overrides import override
from typing_extensions import Annotated
class Connection:
"""A threadpool connection that returns itself to the pool on close()"""
_pool: "Pool"
_db_file: str
_conn: sqlite3.Connection
def __init__(
self, pool: "Pool", db_file: str, is_uri: bool, *args: Any, **kwargs: Any
):
self._pool = pool
self._db_file = db_file
self._conn = sqlite3.connect(
db_file, timeout=1000, check_same_thread=False, uri=is_uri, *args, **kwargs
) # type: ignore
self._conn.isolation_level = None # Handle commits explicitly
def execute(self, sql: str, parameters=...) -> sqlite3.Cursor: # type: ignore
if parameters is ...:
return self._conn.execute(sql)
return self._conn.execute(sql, parameters)
def commit(self) -> None:
self._conn.commit()
def rollback(self) -> None:
self._conn.rollback()
def cursor(self) -> sqlite3.Cursor:
return self._conn.cursor()
def close_actual(self) -> None:
"""Actually closes the connection to the db"""
self._conn.close()
class Pool(ABC):
"""Abstract base class for a pool of connections to a sqlite database."""
@abstractmethod
def __init__(self, db_file: str, is_uri: bool) -> None:
pass
@abstractmethod
def connect(self, *args: Any, **kwargs: Any) -> Connection:
"""Return a connection from the pool."""
pass
@abstractmethod
def close(self) -> None:
"""Close all connections in the pool."""
pass
@abstractmethod
def return_to_pool(self, conn: Connection) -> None:
"""Return a connection to the pool."""
pass
class LockPool(Pool):
"""A pool that has a single connection per thread but uses a lock to ensure that only one thread can use it at a time.
This is used because sqlite does not support multithreaded access with connection timeouts when using the
shared cache mode. We use the shared cache mode to allow multiple threads to share a database.
"""
_connections: Set[Annotated[weakref.ReferenceType, Connection]]
_lock: threading.RLock
_connection: threading.local
_db_file: str
_is_uri: bool
def __init__(self, db_file: str, is_uri: bool = False):
self._connections = set()
self._connection = threading.local()
self._lock = threading.RLock()
self._db_file = db_file
self._is_uri = is_uri
@override
def connect(self, *args: Any, **kwargs: Any) -> Connection:
self._lock.acquire()
if hasattr(self._connection, "conn") and self._connection.conn is not None:
return self._connection.conn # type: ignore # cast doesn't work here for some reason
else:
new_connection = Connection(
self, self._db_file, self._is_uri, *args, **kwargs
)
self._connection.conn = new_connection
self._connections.add(weakref.ref(new_connection))
return new_connection
@override
def return_to_pool(self, conn: Connection) -> None:
try:
self._lock.release()
except RuntimeError:
pass
@override
def close(self) -> None:
for conn in self._connections:
if conn() is not None:
conn().close_actual() # type: ignore
self._connections.clear()
self._connection = threading.local()
try:
self._lock.release()
except RuntimeError:
pass
class PerThreadPool(Pool):
"""Maintains a connection per thread. For now this does not maintain a cap on the number of connections, but it could be
extended to do so and block on connect() if the cap is reached.
"""
_connections: Set[Annotated[weakref.ReferenceType, Connection]]
_lock: threading.Lock
_connection: threading.local
_db_file: str
_is_uri_: bool
def __init__(self, db_file: str, is_uri: bool = False):
self._connections = set()
self._connection = threading.local()
self._lock = threading.Lock()
self._db_file = db_file
self._is_uri = is_uri
@override
def connect(self, *args: Any, **kwargs: Any) -> Connection:
if hasattr(self._connection, "conn") and self._connection.conn is not None:
return self._connection.conn # type: ignore # cast doesn't work here for some reason
else:
new_connection = Connection(
self, self._db_file, self._is_uri, *args, **kwargs
)
self._connection.conn = new_connection
with self._lock:
self._connections.add(weakref.ref(new_connection))
return new_connection
@override
def close(self) -> None:
with self._lock:
for conn in self._connections:
if conn() is not None:
conn().close_actual() # type: ignore
self._connections.clear()
self._connection = threading.local()
@override
def return_to_pool(self, conn: Connection) -> None:
pass # Each thread gets its own connection, so we don't need to return it to the pool
|