meerschaum.connectors

290def make_connector(cls, _is_executor: bool = False):
291    """
292    Register a class as a `Connector`.
293    The `type` will be the lower case of the class name, without the suffix `connector`.
294
295    Parameters
296    ----------
297    instance: bool, default False
298        If `True`, make this connector type an instance connector.
299        This requires implementing the various pipes functions and lots of testing.
300
301    Examples
302    --------
303    >>> import meerschaum as mrsm
304    >>> from meerschaum.connectors import make_connector, Connector
305    >>> 
306    >>> @make_connector
307    >>> class FooConnector(Connector):
308    ...     REQUIRED_ATTRIBUTES: list[str] = ['username', 'password']
309    ... 
310    >>> conn = mrsm.get_connector('foo:bar', username='dog', password='cat')
311    >>> print(conn.username, conn.password)
312    dog cat
313    >>> 
314    """
315    import re
316    suffix_regex = (
317        r'connector$'
318        if not _is_executor
319        else r'executor$'
320    )
321    typ = re.sub(suffix_regex, '', cls.__name__.lower())
322    with _locks['types']:
323        types[typ] = cls
324    with _locks['custom_types']:
325        custom_types.add(typ)
326    with _locks['connectors']:
327        if typ not in connectors:
328            connectors[typ] = {}
329    if getattr(cls, 'IS_INSTANCE', False):
330        with _locks['instance_types']:
331            if typ not in instance_types:
332                instance_types.append(typ)
333
334    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
 58        ### NOTE: Override `REQUIRED_ATTRIBUTES` if `uri` is set.
 59        self.verify_attributes(
 60            ['uri']
 61            if 'uri' in self.__dict__
 62            else getattr(self, 'REQUIRED_ATTRIBUTES', None)
 63        )
 64
 65    def _reset_attributes(self):
 66        self.__dict__ = self._original_dict
 67
 68    def _set_attributes(
 69        self,
 70        *args,
 71        inherit_default: bool = True,
 72        **kw: Any
 73    ):
 74        from meerschaum.config.static import STATIC_CONFIG
 75        from meerschaum.utils.warnings import error
 76
 77        self._attributes = {}
 78
 79        default_label = STATIC_CONFIG['connectors']['default_label']
 80
 81        ### NOTE: Support the legacy method of explicitly passing the type.
 82        label = kw.get('label', None)
 83        if label is None:
 84            if len(args) == 2:
 85                label = args[1]
 86            elif len(args) == 0:
 87                label = None
 88            else:
 89                label = args[0]
 90
 91        if label == 'default':
 92            error(
 93                f"Label cannot be 'default'. Did you mean '{default_label}'?",
 94                InvalidAttributesError,
 95            )
 96        self.__dict__['label'] = label
 97
 98        from meerschaum.config import get_config
 99        conn_configs = copy.deepcopy(get_config('meerschaum', 'connectors'))
100        connector_config = copy.deepcopy(get_config('system', 'connectors'))
101
102        ### inherit attributes from 'default' if exists
103        if inherit_default:
104            inherit_from = 'default'
105            if self.type in conn_configs and inherit_from in conn_configs[self.type]:
106                _inherit_dict = copy.deepcopy(conn_configs[self.type][inherit_from])
107                self._attributes.update(_inherit_dict)
108
109        ### load user config into self._attributes
110        if self.type in conn_configs and self.label in conn_configs[self.type]:
111            self._attributes.update(conn_configs[self.type][self.label] or {})
112
113        ### load system config into self._sys_config
114        ### (deep copy so future Connectors don't inherit changes)
115        if self.type in connector_config:
116            self._sys_config = copy.deepcopy(connector_config[self.type])
117
118        ### add additional arguments or override configuration
119        self._attributes.update(kw)
120
121        ### finally, update __dict__ with _attributes.
122        self.__dict__.update(self._attributes)
123
124    def verify_attributes(
125        self,
126        required_attributes: Optional[List[str]] = None,
127        debug: bool = False,
128    ) -> None:
129        """
130        Ensure that the required attributes have been met.
131        
132        The Connector base class checks the minimum requirements.
133        Child classes may enforce additional requirements.
134
135        Parameters
136        ----------
137        required_attributes: Optional[List[str]], default None
138            Attributes to be verified. If `None`, default to `['label']`.
139
140        debug: bool, default False
141            Verbosity toggle.
142
143        Returns
144        -------
145        Don't return anything.
146
147        Raises
148        ------
149        An error if any of the required attributes are missing.
150        """
151        from meerschaum.utils.warnings import error, warn
152        from meerschaum.utils.debug import dprint
153        from meerschaum.utils.misc import items_str
154        if required_attributes is None:
155            required_attributes = ['label']
156
157        missing_attributes = set()
158        for a in required_attributes:
159            if a not in self.__dict__:
160                missing_attributes.add(a)
161        if len(missing_attributes) > 0:
162            error(
163                (
164                    f"Missing {items_str(list(missing_attributes))} "
165                    + f"for connector '{self.type}:{self.label}'."
166                ),
167                InvalidAttributesError,
168                silent=True,
169                stack=False
170            )
171
172
173    def __str__(self):
174        """
175        When cast to a string, return type:label.
176        """
177        return f"{self.type}:{self.label}"
178
179    def __repr__(self):
180        """
181        Represent the connector as type:label.
182        """
183        return str(self)
184
185    @property
186    def meta(self) -> Dict[str, Any]:
187        """
188        Return the keys needed to reconstruct this Connector.
189        """
190        _meta = {
191            key: value
192            for key, value in self.__dict__.items()
193            if not str(key).startswith('_')
194        }
195        _meta.update({
196            'type': self.type,
197            'label': self.label,
198        })
199        return _meta
200
201
202    @property
203    def type(self) -> str:
204        """
205        Return the type for this connector.
206        """
207        _type = self.__dict__.get('type', None)
208        if _type is None:
209            import re
210            is_executor = self.__class__.__name__.lower().endswith('executor')
211            suffix_regex = (
212                r'connector$'
213                if not is_executor
214                else r'executor$'
215            )
216            _type = re.sub(suffix_regex, '', self.__class__.__name__.lower())
217            self.__dict__['type'] = _type
218        return _type
219
220
221    @property
222    def label(self) -> str:
223        """
224        Return the label for this connector.
225        """
226        _label = self.__dict__.get('label', None)
227        if _label is None:
228            from meerschaum.config.static import STATIC_CONFIG
229            _label = STATIC_CONFIG['connectors']['default_label']
230            self.__dict__['label'] = _label
231        return _label

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