meerschaum.connectors

 20class Connector(metaclass=abc.ABCMeta):
 21    """
 22    The base connector class to hold connection attributes,
 23    """
 24    def __init__(
 25            self,
 26            type: Optional[str] = None,
 27            label: Optional[str] = None,
 28            **kw: Any
 29        ):
 30        """
 31        Parameters
 32        ----------
 33        type: str
 34            The `type` of the connector (e.g. `sql`, `api`, `plugin`).
 35
 36        label: str
 37            The `label` for the connector.
 38
 39        Run `mrsm edit config` and to edit connectors in the YAML file:
 40
 41        ```
 42        meerschaum:
 43            connections:
 44                {type}:
 45                    {label}:
 46                        ### attributes go here
 47        ```
 48
 49        """
 50        self._original_dict = copy.deepcopy(self.__dict__)
 51        self._set_attributes(type=type, label=label, **kw)
 52        self.verify_attributes(getattr(self, 'REQUIRED_ATTRIBUTES', None))
 53
 54    def _reset_attributes(self):
 55        self.__dict__ = self._original_dict
 56
 57    def _set_attributes(
 58            self,
 59            *args,
 60            inherit_default: bool = True,
 61            **kw: Any
 62        ):
 63        from meerschaum.config.static import STATIC_CONFIG
 64        from meerschaum.utils.warnings import error
 65
 66        self._attributes = {}
 67
 68        default_label = STATIC_CONFIG['connectors']['default_label']
 69
 70        ### NOTE: Support the legacy method of explicitly passing the type.
 71        label = kw.get('label', None)
 72        if label is None:
 73            if len(args) == 2:
 74                label = args[1]
 75            elif len(args) == 0:
 76                label = None
 77            else:
 78                label = args[0]
 79
 80        if label == 'default':
 81            error(
 82                f"Label cannot be 'default'. Did you mean '{default_label}'?",
 83                InvalidAttributesError,
 84            )
 85        self.__dict__['label'] = label
 86
 87        from meerschaum.config import get_config
 88        conn_configs = copy.deepcopy(get_config('meerschaum', 'connectors'))
 89        connector_config = copy.deepcopy(get_config('system', 'connectors'))
 90
 91        ### inherit attributes from 'default' if exists
 92        if inherit_default:
 93            inherit_from = 'default'
 94            if self.type in conn_configs and inherit_from in conn_configs[self.type]:
 95                _inherit_dict = copy.deepcopy(conn_configs[self.type][inherit_from])
 96                self._attributes.update(_inherit_dict)
 97
 98        ### load user config into self._attributes
 99        if self.type in conn_configs and self.label in conn_configs[self.type]:
100            self._attributes.update(conn_configs[self.type][self.label])
101
102        ### load system config into self._sys_config
103        ### (deep copy so future Connectors don't inherit changes)
104        if self.type in connector_config:
105            self._sys_config = copy.deepcopy(connector_config[self.type])
106
107        ### add additional arguments or override configuration
108        self._attributes.update(kw)
109
110        ### finally, update __dict__ with _attributes.
111        self.__dict__.update(self._attributes)
112
113
114    def verify_attributes(
115            self,
116            required_attributes: Optional[List[str]] = None,
117            debug: bool = False
118        ) -> None:
119        """
120        Ensure that the required attributes have been met.
121        
122        The Connector base class checks the minimum requirements.
123        Child classes may enforce additional requirements.
124
125        Parameters
126        ----------
127        required_attributes: Optional[List[str]], default None
128            Attributes to be verified. If `None`, default to `['label']`.
129
130        debug: bool, default False
131            Verbosity toggle.
132
133        Returns
134        -------
135        Don't return anything.
136
137        Raises
138        ------
139        An error if any of the required attributes are missing.
140        """
141        from meerschaum.utils.warnings import error, warn
142        from meerschaum.utils.debug import dprint
143        from meerschaum.utils.misc import items_str
144        if required_attributes is None:
145            required_attributes = ['label']
146        missing_attributes = set()
147        for a in required_attributes:
148            if a not in self.__dict__:
149                missing_attributes.add(a)
150        if len(missing_attributes) > 0:
151            error(
152                (
153                    f"Missing {items_str(list(missing_attributes))} "
154                    + f"for connector '{self.type}:{self.label}'."
155                ),
156                InvalidAttributesError,
157                silent = True,
158                stack = False
159            )
160
161
162    def __str__(self):
163        """
164        When cast to a string, return type:label.
165        """
166        return f"{self.type}:{self.label}"
167
168    def __repr__(self):
169        """
170        Represent the connector as type:label.
171        """
172        return str(self)
173
174    @property
175    def meta(self) -> Dict[str, Any]:
176        """
177        Return the keys needed to reconstruct this Connector.
178        """
179        _meta = {
180            key: value
181            for key, value in self.__dict__.items()
182            if not str(key).startswith('_')
183        }
184        _meta.update({
185            'type': self.type,
186            'label': self.label,
187        })
188        return _meta
189
190
191    @property
192    def type(self) -> str:
193        """
194        Return the type for this connector.
195        """
196        _type = self.__dict__.get('type', None)
197        if _type is None:
198            import re
199            _type = re.sub(r'connector$', '', self.__class__.__name__.lower())
200            self.__dict__['type'] = _type
201        return _type
202
203
204    @property
205    def label(self) -> str:
206        """
207        Return the label for this connector.
208        """
209        _label = self.__dict__.get('label', None)
210        if _label is None:
211            from meerschaum.config.static import STATIC_CONFIG
212            _label = STATIC_CONFIG['connectors']['default_label']
213            self.__dict__['label'] = _label
214        return _label

 17class SQLConnector(Connector):
 18    """
 19    Connect to SQL databases via `sqlalchemy`.
 20    
 21    SQLConnectors may be used as Meerschaum instance connectors.
 22    Read more about connectors and instances at
 23    https://meerschaum.io/reference/connectors/
 24
 25    """
 26
 27    IS_INSTANCE: bool = True
 28
 29    from ._create_engine import flavor_configs, create_engine
 30    from ._sql import read, value, exec, execute, to_sql, exec_queries
 31    from meerschaum.utils.sql import test_connection
 32    from ._fetch import fetch, get_pipe_metadef
 33    from ._cli import cli, _cli_exit
 34    from ._pipes import (
 35        fetch_pipes_keys,
 36        create_indices,
 37        drop_indices,
 38        get_create_index_queries,
 39        get_drop_index_queries,
 40        get_add_columns_queries,
 41        get_alter_columns_queries,
 42        delete_pipe,
 43        get_pipe_data,
 44        get_pipe_data_query,
 45        register_pipe,
 46        edit_pipe,
 47        get_pipe_id,
 48        get_pipe_attributes,
 49        sync_pipe,
 50        sync_pipe_inplace,
 51        get_sync_time,
 52        pipe_exists,
 53        get_pipe_rowcount,
 54        drop_pipe,
 55        clear_pipe,
 56        deduplicate_pipe,
 57        get_pipe_table,
 58        get_pipe_columns_types,
 59        get_to_sql_dtype,
 60        get_pipe_schema,
 61    )
 62    from ._plugins import (
 63        register_plugin,
 64        delete_plugin,
 65        get_plugin_id,
 66        get_plugin_version,
 67        get_plugins,
 68        get_plugin_user_id,
 69        get_plugin_username,
 70        get_plugin_attributes,
 71    )
 72    from ._users import (
 73        register_user,
 74        get_user_id,
 75        get_users,
 76        edit_user,
 77        delete_user,
 78        get_user_password_hash,
 79        get_user_type,
 80        get_user_attributes,
 81    )
 82    from ._uri import from_uri, parse_uri
 83    from ._instance import (
 84        _log_temporary_tables_creation,
 85        _drop_temporary_table,
 86        _drop_temporary_tables,
 87        _drop_old_temporary_tables,
 88    )
 89    
 90    def __init__(
 91        self,
 92        label: Optional[str] = None,
 93        flavor: Optional[str] = None,
 94        wait: bool = False,
 95        connect: bool = False,
 96        debug: bool = False,
 97        **kw: Any
 98    ):
 99        """
100        Parameters
101        ----------
102        label: str, default 'main'
103            The identifying label for the connector.
104            E.g. for `sql:main`, 'main' is the label.
105            Defaults to 'main'.
106
107        flavor: Optional[str], default None
108            The database flavor, e.g.
109            `'sqlite'`, `'postgresql'`, `'cockroachdb'`, etc.
110            To see supported flavors, run the `bootstrap connectors` command.
111
112        wait: bool, default False
113            If `True`, block until a database connection has been made.
114            Defaults to `False`.
115
116        connect: bool, default False
117            If `True`, immediately attempt to connect the database and raise
118            a warning if the connection fails.
119            Defaults to `False`.
120
121        debug: bool, default False
122            Verbosity toggle.
123            Defaults to `False`.
124
125        kw: Any
126            All other arguments will be passed to the connector's attributes.
127            Therefore, a connector may be made without being registered,
128            as long enough parameters are supplied to the constructor.
129        """
130        if 'uri' in kw:
131            uri = kw['uri']
132            if uri.startswith('postgres://'):
133                uri = uri.replace('postgres://', 'postgresql://', 1)
134            if uri.startswith('timescaledb://'):
135                uri = uri.replace('timescaledb://', 'postgresql://', 1)
136                flavor = 'timescaledb'
137            kw['uri'] = uri
138            from_uri_params = self.from_uri(kw['uri'], as_dict=True)
139            label = label or from_uri_params.get('label', None)
140            _ = from_uri_params.pop('label', None)
141
142            ### Sometimes the flavor may be provided with a URI.
143            kw.update(from_uri_params)
144            if flavor:
145                kw['flavor'] = flavor
146
147
148        ### set __dict__ in base class
149        super().__init__(
150            'sql',
151            label = label or self.__dict__.get('label', None),
152            **kw
153        )
154
155        if self.__dict__.get('flavor', None) == 'sqlite':
156            self._reset_attributes()
157            self._set_attributes(
158                'sql',
159                label = label,
160                inherit_default = False,
161                **kw
162            )
163            ### For backwards compatability reasons, set the path for sql:local if its missing.
164            if self.label == 'local' and not self.__dict__.get('database', None):
165                from meerschaum.config._paths import SQLITE_DB_PATH
166                self.database = str(SQLITE_DB_PATH)
167
168        ### ensure flavor and label are set accordingly
169        if 'flavor' not in self.__dict__:
170            if flavor is None and 'uri' not in self.__dict__:
171                raise Exception(
172                    f"    Missing flavor. Provide flavor as a key for '{self}'."
173                )
174            self.flavor = flavor or self.parse_uri(self.__dict__['uri']).get('flavor', None)
175
176        if self.flavor == 'postgres':
177            self.flavor = 'postgresql'
178
179        self._debug = debug
180        ### Store the PID and thread at initialization
181        ### so we can dispose of the Pool in child processes or threads.
182        import os, threading
183        self._pid = os.getpid()
184        self._thread_ident = threading.current_thread().ident
185        self._sessions = {}
186        self._locks = {'_sessions': threading.RLock(), }
187
188        ### verify the flavor's requirements are met
189        if self.flavor not in self.flavor_configs:
190            error(f"Flavor '{self.flavor}' is not supported by Meerschaum SQLConnector")
191        if not self.__dict__.get('uri'):
192            self.verify_attributes(
193                self.flavor_configs[self.flavor].get('requirements', set()),
194                debug=debug,
195            )
196
197        if wait:
198            from meerschaum.connectors.poll import retry_connect
199            retry_connect(connector=self, debug=debug)
200
201        if connect:
202            if not self.test_connection(debug=debug):
203                from meerschaum.utils.warnings import warn
204                warn(f"Failed to connect with connector '{self}'!", stack=False)
205
206    @property
207    def Session(self):
208        if '_Session' not in self.__dict__:
209            if self.engine is None:
210                return None
211
212            from meerschaum.utils.packages import attempt_import
213            sqlalchemy_orm = attempt_import('sqlalchemy.orm')
214            session_factory = sqlalchemy_orm.sessionmaker(self.engine)
215            self._Session = sqlalchemy_orm.scoped_session(session_factory)
216
217        return self._Session
218
219    @property
220    def engine(self):
221        import os, threading
222        ### build the sqlalchemy engine
223        if '_engine' not in self.__dict__:
224            self._engine, self._engine_str = self.create_engine(include_uri=True)
225
226        same_process = os.getpid() == self._pid
227        same_thread = threading.current_thread().ident == self._thread_ident
228
229        ### handle child processes
230        if not same_process:
231            self._pid = os.getpid()
232            self._thread = threading.current_thread()
233            from meerschaum.utils.warnings import warn
234            warn(f"Different PID detected. Disposing of connections...")
235            self._engine.dispose()
236
237        ### handle different threads
238        if not same_thread:
239            pass
240
241        return self._engine
242
243    @property
244    def DATABASE_URL(self) -> str:
245        """
246        Return the URI connection string (alias for `SQLConnector.URI`.
247        """
248        _ = self.engine
249        return str(self._engine_str)
250
251    @property
252    def URI(self) -> str:
253        """
254        Return the URI connection string.
255        """
256        _ = self.engine
257        return str(self._engine_str)
258
259    @property
260    def IS_THREAD_SAFE(self) -> str:
261        """
262        Return whether this connector may be multithreaded.
263        """
264        if self.flavor == 'duckdb':
265            return False
266        if self.flavor == 'sqlite':
267            return ':memory:' not in self.URI
268        return True
269
270
271    @property
272    def metadata(self):
273        """
274        Return the metadata bound to this configured schema.
275        """
276        from meerschaum.utils.packages import attempt_import
277        sqlalchemy = attempt_import('sqlalchemy')
278        if '_metadata' not in self.__dict__:
279            self._metadata = sqlalchemy.MetaData(schema=self.schema)
280        return self._metadata
281
282
283    @property
284    def instance_schema(self):
285        """
286        Return the schema name for Meerschaum tables. 
287        """
288        return self.schema
289
290
291    @property
292    def internal_schema(self):
293        """
294        Return the schema name for internal tables. 
295        """
296        from meerschaum.config.static import STATIC_CONFIG
297        from meerschaum.utils.packages import attempt_import
298        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
299        schema_name = self.__dict__.get('internal_schema', None) or (
300            STATIC_CONFIG['sql']['internal_schema']
301            if self.flavor not in NO_SCHEMA_FLAVORS
302            else self.schema
303        )
304
305        if '_internal_schema' not in self.__dict__:
306            self._internal_schema = schema_name
307        return self._internal_schema
308
309
310    @property
311    def db(self) -> Optional[databases.Database]:
312        from meerschaum.utils.packages import attempt_import
313        databases = attempt_import('databases', lazy=False, install=True)
314        url = self.DATABASE_URL
315        if 'mysql' in url:
316            url = url.replace('+pymysql', '')
317        if '_db' not in self.__dict__:
318            try:
319                self._db = databases.Database(url)
320            except KeyError:
321                ### Likely encountered an unsupported flavor.
322                from meerschaum.utils.warnings import warn
323                self._db = None
324        return self._db
325
326
327    @property
328    def db_version(self) -> Union[str, None]:
329        """
330        Return the database version.
331        """
332        _db_version = self.__dict__.get('_db_version', None)
333        if _db_version is not None:
334            return _db_version
335
336        from meerschaum.utils.sql import get_db_version
337        self._db_version = get_db_version(self)
338        return self._db_version
339
340
341    @property
342    def schema(self) -> Union[str, None]:
343        """
344        Return the default schema to use.
345        A value of `None` will not prepend a schema.
346        """
347        if 'schema' in self.__dict__:
348            return self.__dict__['schema']
349
350        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
351        if self.flavor in NO_SCHEMA_FLAVORS:
352            self.__dict__['schema'] = None
353            return None
354
355        sqlalchemy = mrsm.attempt_import('sqlalchemy')
356        _schema = sqlalchemy.inspect(self.engine).default_schema_name
357        self.__dict__['schema'] = _schema
358        return _schema
359
360
361    def __getstate__(self):
362        return self.__dict__
363
364    def __setstate__(self, d):
365        self.__dict__.update(d)
366
367    def __call__(self):
368        return self