|
8 | 8 | import sys |
9 | 9 | import typing as t |
10 | 10 | from enum import Enum |
| 11 | +from functools import partial |
11 | 12 |
|
12 | 13 | from pydantic import Field |
13 | 14 | from sqlglot import exp |
@@ -84,15 +85,21 @@ def is_recommended_for_state_sync(self) -> bool: |
84 | 85 | """Whether this connection is recommended for being used as a state sync for production state syncs""" |
85 | 86 | return self.type_ in RECOMMENDED_STATE_SYNC_ENGINES |
86 | 87 |
|
| 88 | + @property |
| 89 | + def _connection_factory_with_kwargs(self) -> t.Callable[[], t.Any]: |
| 90 | + """A function that is called to return a connection object for the given Engine Adapter""" |
| 91 | + return partial( |
| 92 | + self._connection_factory, |
| 93 | + **{ |
| 94 | + **self._static_connection_kwargs, |
| 95 | + **{k: v for k, v in self.dict().items() if k in self._connection_kwargs_keys}, |
| 96 | + }, |
| 97 | + ) |
| 98 | + |
87 | 99 | def create_engine_adapter(self, register_comments_override: bool = False) -> EngineAdapter: |
88 | 100 | """Returns a new instance of the Engine Adapter.""" |
89 | 101 | return self._engine_adapter( |
90 | | - lambda: self._connection_factory( |
91 | | - **{ |
92 | | - **self._static_connection_kwargs, |
93 | | - **{k: v for k, v in self.dict().items() if k in self._connection_kwargs_keys}, |
94 | | - } |
95 | | - ), |
| 102 | + self._connection_factory_with_kwargs, |
96 | 103 | multithreaded=self.concurrent_tasks > 1, |
97 | 104 | cursor_kwargs=self._cursor_kwargs, |
98 | 105 | default_catalog=self.get_catalog(), |
|
0 commit comments