sqlmesh.utils.connection_pool

  1import abc
  2import logging
  3import typing as t
  4from collections import defaultdict
  5from threading import Lock, get_ident
  6
  7logger = logging.getLogger(__name__)
  8
  9
 10class ConnectionPool(abc.ABC):
 11    @abc.abstractmethod
 12    def get_cursor(self) -> t.Any:
 13        """Returns cached cursor instance.
 14
 15        Automatically creates a new instance if one is not available.
 16
 17        Returns:
 18            A cursor instance.
 19        """
 20
 21    @abc.abstractmethod
 22    def get(self) -> t.Any:
 23        """Returns cached connection instance.
 24
 25        Automatically opens a new connection if one is not available.
 26
 27        Returns:
 28            A connection instance.
 29        """
 30
 31    @abc.abstractmethod
 32    def get_attribute(self, key: str) -> t.Optional[t.Any]:
 33        """Returns an attribute associated with the connection.
 34
 35        Args:
 36            key: Attribute key.
 37
 38        Returns:
 39            Attribute value or None if not found.
 40        """
 41
 42    @abc.abstractmethod
 43    def set_attribute(self, key: str, value: t.Any) -> None:
 44        """Sets an attribute associated with the connection.
 45
 46        Args:
 47            key: Attribute key.
 48            value: Attribute value.
 49        """
 50
 51    @abc.abstractmethod
 52    def get_all_attributes(self, key: str) -> t.List[t.Any]:
 53        """Returns all attributes with the given key across all connections/threads.
 54
 55        Args:
 56            key: Attribute key.
 57
 58        Returns:
 59            List of attribute values from all connections/threads.
 60        """
 61
 62    @abc.abstractmethod
 63    def begin(self) -> None:
 64        """Starts a new transaction."""
 65
 66    @abc.abstractmethod
 67    def commit(self) -> None:
 68        """Commits the current transaction."""
 69
 70    @abc.abstractmethod
 71    def rollback(self) -> None:
 72        """Rolls back the current transaction."""
 73
 74    @property
 75    @abc.abstractmethod
 76    def is_transaction_active(self) -> bool:
 77        """Returns True if there is an active transaction and False otherwise."""
 78
 79    @abc.abstractmethod
 80    def close_cursor(self) -> None:
 81        """Closes the current cursor instance if exists."""
 82
 83    @abc.abstractmethod
 84    def close(self) -> None:
 85        """Closes the current connection instance if exists.
 86
 87        Note: if there is a cursor instance available it will be closed as well.
 88        """
 89
 90    @abc.abstractmethod
 91    def close_all(self, exclude_calling_thread: bool = False) -> None:
 92        """Closes all cached cursors and connections.
 93
 94        Args:
 95            exclude_calling_thread: If set to True excludes cursors and connections associated
 96                with the calling thread.
 97        """
 98
 99
100class _TransactionManagementMixin(ConnectionPool):
101    def _do_begin(self) -> None:
102        cursor = self.get_cursor()
103        if hasattr(cursor, "begin"):
104            cursor.begin()
105        else:
106            conn = self.get()
107            if hasattr(conn, "begin"):
108                conn.begin()
109
110    def _do_commit(self) -> None:
111        cursor = self.get_cursor()
112        if hasattr(cursor, "commit"):
113            cursor.commit()
114        else:
115            self.get().commit()
116
117    def _do_rollback(self) -> None:
118        cursor = self.get_cursor()
119        if hasattr(cursor, "rollback"):
120            cursor.rollback()
121        else:
122            self.get().rollback()
123
124
125class _ThreadLocalBase(_TransactionManagementMixin):
126    def __init__(
127        self,
128        connection_factory: t.Callable[[], t.Any],
129        cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
130    ):
131        self._connection_factory = connection_factory
132        self._thread_cursors: t.Dict[t.Hashable, t.Any] = {}
133        self._thread_transactions: t.Set[t.Hashable] = set()
134        self._thread_attributes: t.Dict[t.Hashable, t.Dict[str, t.Any]] = defaultdict(dict)
135        self._thread_cursors_lock = Lock()
136        self._thread_transactions_lock = Lock()
137        self._cursor_init = cursor_init
138
139    def get_cursor(self) -> t.Any:
140        thread_id = get_ident()
141        with self._thread_cursors_lock:
142            if thread_id not in self._thread_cursors:
143                self._thread_cursors[thread_id] = self.get().cursor()
144                if self._cursor_init:
145                    self._cursor_init(self._thread_cursors[thread_id])
146            return self._thread_cursors[thread_id]
147
148    def get_attribute(self, key: str) -> t.Optional[t.Any]:
149        thread_id = get_ident()
150        return self._thread_attributes[thread_id].get(key)
151
152    def set_attribute(self, key: str, value: t.Any) -> None:
153        thread_id = get_ident()
154        self._thread_attributes[thread_id][key] = value
155
156    def get_all_attributes(self, key: str) -> t.List[t.Any]:
157        """Returns all attributes with the given key across all threads."""
158        return [
159            thread_attrs[key]
160            for thread_attrs in self._thread_attributes.values()
161            if key in thread_attrs
162        ]
163
164    def begin(self) -> None:
165        self._do_begin()
166        with self._thread_transactions_lock:
167            self._thread_transactions.add(get_ident())
168
169    def commit(self) -> None:
170        self._do_commit()
171        self._discard_transaction(get_ident())
172
173    def rollback(self) -> None:
174        self._do_rollback()
175        self._discard_transaction(get_ident())
176
177    @property
178    def is_transaction_active(self) -> bool:
179        with self._thread_transactions_lock:
180            return get_ident() in self._thread_transactions
181
182    def close_cursor(self) -> None:
183        thread_id = get_ident()
184        with self._thread_cursors_lock:
185            if thread_id in self._thread_cursors:
186                _try_close(self._thread_cursors[thread_id], "cursor")
187                self._thread_cursors.pop(thread_id)
188
189    def _discard_transaction(self, thread_id: t.Hashable) -> None:
190        with self._thread_transactions_lock:
191            self._thread_transactions.discard(thread_id)
192
193
194class ThreadLocalConnectionPool(_ThreadLocalBase):
195    def __init__(
196        self,
197        connection_factory: t.Callable[[], t.Any],
198        cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
199    ):
200        super().__init__(connection_factory, cursor_init)
201        self._thread_connections: t.Dict[t.Hashable, t.Any] = {}
202        self._thread_connections_lock = Lock()
203
204    def get(self) -> t.Any:
205        thread_id = get_ident()
206        with self._thread_connections_lock:
207            if thread_id not in self._thread_connections:
208                self._thread_connections[thread_id] = self._connection_factory()
209            return self._thread_connections[thread_id]
210
211    def close(self) -> None:
212        thread_id = get_ident()
213        with self._thread_cursors_lock, self._thread_connections_lock:
214            if thread_id in self._thread_connections:
215                _try_close(self._thread_connections[thread_id], "connection")
216                self._thread_connections.pop(thread_id)
217                self._thread_cursors.pop(thread_id, None)
218                self._discard_transaction(thread_id)
219            self._thread_attributes.pop(thread_id, None)
220
221    def close_all(self, exclude_calling_thread: bool = False) -> None:
222        calling_thread_id = get_ident()
223        with self._thread_cursors_lock, self._thread_connections_lock:
224            for thread_id, connection in self._thread_connections.copy().items():
225                if not exclude_calling_thread or thread_id != calling_thread_id:
226                    _try_close(connection, "connection")
227                    self._thread_connections.pop(thread_id)
228                    self._thread_cursors.pop(thread_id, None)
229                    self._discard_transaction(thread_id)
230
231            self._thread_attributes.clear()
232
233
234class ThreadLocalSharedConnectionPool(_ThreadLocalBase):
235    def __init__(
236        self,
237        connection_factory: t.Callable[[], t.Any],
238        cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
239    ):
240        super().__init__(connection_factory, cursor_init)
241        self._connection: t.Optional[t.Any] = None
242        self._connection_lock = Lock()
243
244    def get(self) -> t.Any:
245        with self._connection_lock:
246            if self._connection is None:
247                self._connection = self._connection_factory()
248            return self._connection
249
250    def close(self) -> None:
251        thread_id = get_ident()
252        with self._thread_cursors_lock, self._connection_lock:
253            if thread_id in self._thread_cursors:
254                _try_close(self._thread_cursors[thread_id], "cursor")
255                self._thread_cursors.pop(thread_id)
256            self._discard_transaction(thread_id)
257            self._thread_attributes.pop(thread_id, None)
258
259    def close_all(self, exclude_calling_thread: bool = False) -> None:
260        calling_thread_id = get_ident()
261        with self._thread_cursors_lock, self._connection_lock:
262            for thread_id, cursor in self._thread_cursors.copy().items():
263                if not exclude_calling_thread or thread_id != calling_thread_id:
264                    _try_close(cursor, "cursor")
265                    self._thread_cursors.pop(thread_id)
266                    self._discard_transaction(thread_id)
267                self._thread_attributes.pop(thread_id, None)
268
269            if not exclude_calling_thread:
270                _try_close(self._connection, "connection")
271                self._connection = None
272
273
274class SingletonConnectionPool(_TransactionManagementMixin):
275    def __init__(
276        self,
277        connection_factory: t.Callable[[], t.Any],
278        cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
279    ):
280        self._connection_factory = connection_factory
281        self._connection: t.Optional[t.Any] = None
282        self._cursor: t.Optional[t.Any] = None
283        self._attributes: t.Dict[str, t.Any] = {}
284        self._is_transaction_active: bool = False
285        self._cursor_init = cursor_init
286
287    def get_cursor(self) -> t.Any:
288        if not self._cursor:
289            self._cursor = self.get().cursor()
290            if self._cursor_init:
291                self._cursor_init(self._cursor)
292        return self._cursor
293
294    def get(self) -> t.Any:
295        if not self._connection:
296            self._connection = self._connection_factory()
297        return self._connection
298
299    def get_attribute(self, key: str) -> t.Optional[t.Any]:
300        return self._attributes.get(key)
301
302    def set_attribute(self, key: str, value: t.Any) -> None:
303        self._attributes[key] = value
304
305    def get_all_attributes(self, key: str) -> t.List[t.Any]:
306        """Returns all attributes with the given key (single-threaded pool has at most one)."""
307        value = self._attributes.get(key)
308        return [value] if value is not None else []
309
310    def begin(self) -> None:
311        self._do_begin()
312        self._is_transaction_active = True
313
314    def commit(self) -> None:
315        self._do_commit()
316        self._is_transaction_active = False
317
318    def rollback(self) -> None:
319        self._do_rollback()
320        self._is_transaction_active = False
321
322    @property
323    def is_transaction_active(self) -> bool:
324        return self._is_transaction_active
325
326    def close_cursor(self) -> None:
327        _try_close(self._cursor, "cursor")
328        self._cursor = None
329
330    def close(self) -> None:
331        _try_close(self._connection, "connection")
332        self._connection = None
333        self._cursor = None
334        self._is_transaction_active = False
335        self._attributes.clear()
336
337    def close_all(self, exclude_calling_thread: bool = False) -> None:
338        if not exclude_calling_thread:
339            self.close()
340
341
342def create_connection_pool(
343    connection_factory: t.Callable[[], t.Any],
344    multithreaded: bool,
345    shared_connection: bool = False,
346    cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
347) -> ConnectionPool:
348    pool_class = (
349        ThreadLocalSharedConnectionPool
350        if multithreaded and shared_connection
351        else ThreadLocalConnectionPool
352        if multithreaded
353        else SingletonConnectionPool
354    )
355    return pool_class(connection_factory, cursor_init=cursor_init)
356
357
358def _try_close(closeable: t.Any, kind: str) -> None:
359    if closeable is None:
360        return
361    try:
362        closeable.close()
363    except Exception:
364        logger.exception("Failed to close %s", kind)