meerschaum.connectors

278def make_connector(
279        cls,
280    ):
281    """
282    Register a class as a `Connector`.
283    The `type` will be the lower case of the class name, without the suffix `connector`.
284
285    Parameters
286    ----------
287    instance: bool, default False
288        If `True`, make this connector type an instance connector.
289        This requires implementing the various pipes functions and lots of testing.
290
291    Examples
292    --------
293    >>> import meerschaum as mrsm
294    >>> from meerschaum.connectors import make_connector, Connector
295    >>> 
296    >>> @make_connector
297    >>> class FooConnector(Connector):
298    ...     REQUIRED_ATTRIBUTES: list[str] = ['username', 'password']
299    ... 
300    >>> conn = mrsm.get_connector('foo:bar', username='dog', password='cat')
301    >>> print(conn.username, conn.password)
302    dog cat
303    >>> 
304    """
305    import re
306    typ = re.sub(r'connector$', '', cls.__name__.lower())
307    with _locks['types']:
308        types[typ] = cls
309    with _locks['custom_types']:
310        custom_types.add(typ)
311    with _locks['connectors']:
312        if typ not in connectors:
313            connectors[typ] = {}
314    if getattr(cls, 'IS_INSTANCE', False):
315        with _locks['instance_types']:
316            if typ not in instance_types:
317                instance_types.append(typ)
318
319    return cls

 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        Set the given keyword arguments as attributes.
 32
 33        Parameters
 34        ----------
 35        type: str
 36            The `type` of the connector (e.g. `sql`, `api`, `plugin`).
 37
 38        label: str
 39            The `label` for the connector.
 40
 41
 42        Examples
 43        --------
 44        Run `mrsm edit config` and to edit connectors in the YAML file:
 45
 46        ```yaml
 47        meerschaum:
 48            connections:
 49                {type}:
 50                    {label}:
 51                        ### attributes go here
 52        ```
 53
 54        """
 55        self._original_dict = copy.deepcopy(self.__dict__)
 56        self._set_attributes(type=type, label=label, **kw)
 57        self.verify_attributes(getattr(self, 'REQUIRED_ATTRIBUTES', None))
 58
 59    def _reset_attributes(self):
 60        self.__dict__ = self._original_dict
 61
 62    def _set_attributes(
 63            self,
 64            *args,
 65            inherit_default: bool = True,
 66            **kw: Any
 67        ):
 68        from meerschaum.config.static import STATIC_CONFIG
 69        from meerschaum.utils.warnings import error
 70
 71        self._attributes = {}
 72
 73        default_label = STATIC_CONFIG['connectors']['default_label']
 74
 75        ### NOTE: Support the legacy method of explicitly passing the type.
 76        label = kw.get('label', None)
 77        if label is None:
 78            if len(args) == 2:
 79                label = args[1]
 80            elif len(args) == 0:
 81                label = None
 82            else:
 83                label = args[0]
 84
 85        if label == 'default':
 86            error(
 87                f"Label cannot be 'default'. Did you mean '{default_label}'?",
 88                InvalidAttributesError,
 89            )
 90        self.__dict__['label'] = label
 91
 92        from meerschaum.config import get_config
 93        conn_configs = copy.deepcopy(get_config('meerschaum', 'connectors'))
 94        connector_config = copy.deepcopy(get_config('system', 'connectors'))
 95
 96        ### inherit attributes from 'default' if exists
 97        if inherit_default:
 98            inherit_from = 'default'
 99            if self.type in conn_configs and inherit_from in conn_configs[self.type]:
100                _inherit_dict = copy.deepcopy(conn_configs[self.type][inherit_from])
101                self._attributes.update(_inherit_dict)
102
103        ### load user config into self._attributes
104        if self.type in conn_configs and self.label in conn_configs[self.type]:
105            self._attributes.update(conn_configs[self.type][self.label])
106
107        ### load system config into self._sys_config
108        ### (deep copy so future Connectors don't inherit changes)
109        if self.type in connector_config:
110            self._sys_config = copy.deepcopy(connector_config[self.type])
111
112        ### add additional arguments or override configuration
113        self._attributes.update(kw)
114
115        ### finally, update __dict__ with _attributes.
116        self.__dict__.update(self._attributes)
117
118
119    def verify_attributes(
120            self,
121            required_attributes: Optional[List[str]] = None,
122            debug: bool = False
123        ) -> None:
124        """
125        Ensure that the required attributes have been met.
126        
127        The Connector base class checks the minimum requirements.
128        Child classes may enforce additional requirements.
129
130        Parameters
131        ----------
132        required_attributes: Optional[List[str]], default None
133            Attributes to be verified. If `None`, default to `['label']`.
134
135        debug: bool, default False
136            Verbosity toggle.
137
138        Returns
139        -------
140        Don't return anything.
141
142        Raises
143        ------
144        An error if any of the required attributes are missing.
145        """
146        from meerschaum.utils.warnings import error, warn
147        from meerschaum.utils.debug import dprint
148        from meerschaum.utils.misc import items_str
149        if required_attributes is None:
150            required_attributes = ['label']
151        missing_attributes = set()
152        for a in required_attributes:
153            if a not in self.__dict__:
154                missing_attributes.add(a)
155        if len(missing_attributes) > 0:
156            error(
157                (
158                    f"Missing {items_str(list(missing_attributes))} "
159                    + f"for connector '{self.type}:{self.label}'."
160                ),
161                InvalidAttributesError,
162                silent = True,
163                stack = False
164            )
165
166
167    def __str__(self):
168        """
169        When cast to a string, return type:label.
170        """
171        return f"{self.type}:{self.label}"
172
173    def __repr__(self):
174        """
175        Represent the connector as type:label.
176        """
177        return str(self)
178
179    @property
180    def meta(self) -> Dict[str, Any]:
181        """
182        Return the keys needed to reconstruct this Connector.
183        """
184        _meta = {
185            key: value
186            for key, value in self.__dict__.items()
187            if not str(key).startswith('_')
188        }
189        _meta.update({
190            'type': self.type,
191            'label': self.label,
192        })
193        return _meta
194
195
196    @property
197    def type(self) -> str:
198        """
199        Return the type for this connector.
200        """
201        _type = self.__dict__.get('type', None)
202        if _type is None:
203            import re
204            _type = re.sub(r'connector$', '', self.__class__.__name__.lower())
205            self.__dict__['type'] = _type
206        return _type
207
208
209    @property
210    def label(self) -> str:
211        """
212        Return the label for this connector.
213        """
214        _label = self.__dict__.get('label', None)
215        if _label is None:
216            from meerschaum.config.static import STATIC_CONFIG
217            _label = STATIC_CONFIG['connectors']['default_label']
218            self.__dict__['label'] = _label
219        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') and not uri.startswith('postgresql'):
133                uri = uri.replace('postgres', 'postgresql', 1)
134            if uri.startswith('postgresql') and not uri.startswith('postgresql+'):
135                uri = uri.replace('postgresql://', 'postgresql+psycopg', 1)
136            if uri.startswith('timescaledb://'):
137                uri = uri.replace('timescaledb://', 'postgresql://', 1)
138                flavor = 'timescaledb'
139            kw['uri'] = uri
140            from_uri_params = self.from_uri(kw['uri'], as_dict=True)
141            label = label or from_uri_params.get('label', None)
142            _ = from_uri_params.pop('label', None)
143
144            ### Sometimes the flavor may be provided with a URI.
145            kw.update(from_uri_params)
146            if flavor:
147                kw['flavor'] = flavor
148
149
150        ### set __dict__ in base class
151        super().__init__(
152            'sql',
153            label = label or self.__dict__.get('label', None),
154            **kw
155        )
156
157        if self.__dict__.get('flavor', None) == 'sqlite':
158            self._reset_attributes()
159            self._set_attributes(
160                'sql',
161                label = label,
162                inherit_default = False,
163                **kw
164            )
165            ### For backwards compatability reasons, set the path for sql:local if its missing.
166            if self.label == 'local' and not self.__dict__.get('database', None):
167                from meerschaum.config._paths import SQLITE_DB_PATH
168                self.database = str(SQLITE_DB_PATH)
169
170        ### ensure flavor and label are set accordingly
171        if 'flavor' not in self.__dict__:
172            if flavor is None and 'uri' not in self.__dict__:
173                raise Exception(
174                    f"    Missing flavor. Provide flavor as a key for '{self}'."
175                )
176            self.flavor = flavor or self.parse_uri(self.__dict__['uri']).get('flavor', None)
177
178        if self.flavor == 'postgres':
179            self.flavor = 'postgresql'
180
181        self._debug = debug
182        ### Store the PID and thread at initialization
183        ### so we can dispose of the Pool in child processes or threads.
184        import os, threading
185        self._pid = os.getpid()
186        self._thread_ident = threading.current_thread().ident
187        self._sessions = {}
188        self._locks = {'_sessions': threading.RLock(), }
189
190        ### verify the flavor's requirements are met
191        if self.flavor not in self.flavor_configs:
192            error(f"Flavor '{self.flavor}' is not supported by Meerschaum SQLConnector")
193        if not self.__dict__.get('uri'):
194            self.verify_attributes(
195                self.flavor_configs[self.flavor].get('requirements', set()),
196                debug=debug,
197            )
198
199        if wait:
200            from meerschaum.connectors.poll import retry_connect
201            retry_connect(connector=self, debug=debug)
202
203        if connect:
204            if not self.test_connection(debug=debug):
205                from meerschaum.utils.warnings import warn
206                warn(f"Failed to connect with connector '{self}'!", stack=False)
207
208    @property
209    def Session(self):
210        if '_Session' not in self.__dict__:
211            if self.engine is None:
212                return None
213
214            from meerschaum.utils.packages import attempt_import
215            sqlalchemy_orm = attempt_import('sqlalchemy.orm')
216            session_factory = sqlalchemy_orm.sessionmaker(self.engine)
217            self._Session = sqlalchemy_orm.scoped_session(session_factory)
218
219        return self._Session
220
221    @property
222    def engine(self):
223        import os, threading
224        ### build the sqlalchemy engine
225        if '_engine' not in self.__dict__:
226            self._engine, self._engine_str = self.create_engine(include_uri=True)
227
228        same_process = os.getpid() == self._pid
229        same_thread = threading.current_thread().ident == self._thread_ident
230
231        ### handle child processes
232        if not same_process:
233            self._pid = os.getpid()
234            self._thread = threading.current_thread()
235            from meerschaum.utils.warnings import warn
236            warn(f"Different PID detected. Disposing of connections...")
237            self._engine.dispose()
238
239        ### handle different threads
240        if not same_thread:
241            pass
242
243        return self._engine
244
245    @property
246    def DATABASE_URL(self) -> str:
247        """
248        Return the URI connection string (alias for `SQLConnector.URI`.
249        """
250        _ = self.engine
251        return str(self._engine_str)
252
253    @property
254    def URI(self) -> str:
255        """
256        Return the URI connection string.
257        """
258        _ = self.engine
259        return str(self._engine_str)
260
261    @property
262    def IS_THREAD_SAFE(self) -> str:
263        """
264        Return whether this connector may be multithreaded.
265        """
266        if self.flavor == 'duckdb':
267            return False
268        if self.flavor == 'sqlite':
269            return ':memory:' not in self.URI
270        return True
271
272
273    @property
274    def metadata(self):
275        """
276        Return the metadata bound to this configured schema.
277        """
278        from meerschaum.utils.packages import attempt_import
279        sqlalchemy = attempt_import('sqlalchemy')
280        if '_metadata' not in self.__dict__:
281            self._metadata = sqlalchemy.MetaData(schema=self.schema)
282        return self._metadata
283
284
285    @property
286    def instance_schema(self):
287        """
288        Return the schema name for Meerschaum tables. 
289        """
290        return self.schema
291
292
293    @property
294    def internal_schema(self):
295        """
296        Return the schema name for internal tables. 
297        """
298        from meerschaum.config.static import STATIC_CONFIG
299        from meerschaum.utils.packages import attempt_import
300        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
301        schema_name = self.__dict__.get('internal_schema', None) or (
302            STATIC_CONFIG['sql']['internal_schema']
303            if self.flavor not in NO_SCHEMA_FLAVORS
304            else self.schema
305        )
306
307        if '_internal_schema' not in self.__dict__:
308            self._internal_schema = schema_name
309        return self._internal_schema
310
311
312    @property
313    def db(self) -> Optional[databases.Database]:
314        from meerschaum.utils.packages import attempt_import
315        databases = attempt_import('databases', lazy=False, install=True)
316        url = self.DATABASE_URL
317        if 'mysql' in url:
318            url = url.replace('+pymysql', '')
319        if '_db' not in self.__dict__:
320            try:
321                self._db = databases.Database(url)
322            except KeyError:
323                ### Likely encountered an unsupported flavor.
324                from meerschaum.utils.warnings import warn
325                self._db = None
326        return self._db
327
328
329    @property
330    def db_version(self) -> Union[str, None]:
331        """
332        Return the database version.
333        """
334        _db_version = self.__dict__.get('_db_version', None)
335        if _db_version is not None:
336            return _db_version
337
338        from meerschaum.utils.sql import get_db_version
339        self._db_version = get_db_version(self)
340        return self._db_version
341
342
343    @property
344    def schema(self) -> Union[str, None]:
345        """
346        Return the default schema to use.
347        A value of `None` will not prepend a schema.
348        """
349        if 'schema' in self.__dict__:
350            return self.__dict__['schema']
351
352        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
353        if self.flavor in NO_SCHEMA_FLAVORS:
354            self.__dict__['schema'] = None
355            return None
356
357        sqlalchemy = mrsm.attempt_import('sqlalchemy')
358        _schema = sqlalchemy.inspect(self.engine).default_schema_name
359        self.__dict__['schema'] = _schema
360        return _schema
361
362
363    def __getstate__(self):
364        return self.__dict__
365
366    def __setstate__(self, d):
367        self.__dict__.update(d)
368
369    def __call__(self):
370        return self