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        get_temporary_target,
 74        create_pipe_indices,
 75        drop_pipe_indices,
 76        get_pipe_index_names,
 77    )
 78    from ._plugins import (
 79        register_plugin,
 80        delete_plugin,
 81        get_plugin_id,
 82        get_plugin_version,
 83        get_plugins,
 84        get_plugin_user_id,
 85        get_plugin_username,
 86        get_plugin_attributes,
 87    )
 88    from ._users import (
 89        register_user,
 90        get_user_id,
 91        get_users,
 92        edit_user,
 93        delete_user,
 94        get_user_password_hash,
 95        get_user_type,
 96        get_user_attributes,
 97    )
 98    from ._uri import from_uri, parse_uri
 99    from ._instance import (
100        _log_temporary_tables_creation,
101        _drop_temporary_table,
102        _drop_temporary_tables,
103        _drop_old_temporary_tables,
104    )
105
106    def __init__(
107        self,
108        label: Optional[str] = None,
109        flavor: Optional[str] = None,
110        wait: bool = False,
111        connect: bool = False,
112        debug: bool = False,
113        **kw: Any
114    ):
115        """
116        Parameters
117        ----------
118        label: str, default 'main'
119            The identifying label for the connector.
120            E.g. for `sql:main`, 'main' is the label.
121            Defaults to 'main'.
122
123        flavor: Optional[str], default None
124            The database flavor, e.g.
125            `'sqlite'`, `'postgresql'`, `'cockroachdb'`, etc.
126            To see supported flavors, run the `bootstrap connectors` command.
127
128        wait: bool, default False
129            If `True`, block until a database connection has been made.
130            Defaults to `False`.
131
132        connect: bool, default False
133            If `True`, immediately attempt to connect the database and raise
134            a warning if the connection fails.
135            Defaults to `False`.
136
137        debug: bool, default False
138            Verbosity toggle.
139            Defaults to `False`.
140
141        kw: Any
142            All other arguments will be passed to the connector's attributes.
143            Therefore, a connector may be made without being registered,
144            as long enough parameters are supplied to the constructor.
145        """
146        if 'uri' in kw:
147            uri = kw['uri']
148            if uri.startswith('postgres') and not uri.startswith('postgresql'):
149                uri = uri.replace('postgres', 'postgresql', 1)
150            if uri.startswith('postgresql') and not uri.startswith('postgresql+'):
151                uri = uri.replace('postgresql://', 'postgresql+psycopg://', 1)
152            if uri.startswith('timescaledb://'):
153                uri = uri.replace('timescaledb://', 'postgresql+psycopg://', 1)
154                flavor = 'timescaledb'
155            kw['uri'] = uri
156            from_uri_params = self.from_uri(kw['uri'], as_dict=True)
157            label = label or from_uri_params.get('label', None)
158            _ = from_uri_params.pop('label', None)
159
160            ### Sometimes the flavor may be provided with a URI.
161            kw.update(from_uri_params)
162            if flavor:
163                kw['flavor'] = flavor
164
165        ### set __dict__ in base class
166        super().__init__(
167            'sql',
168            label = label or self.__dict__.get('label', None),
169            **kw
170        )
171
172        if self.__dict__.get('flavor', None) == 'sqlite':
173            self._reset_attributes()
174            self._set_attributes(
175                'sql',
176                label = label,
177                inherit_default = False,
178                **kw
179            )
180            ### For backwards compatability reasons, set the path for sql:local if its missing.
181            if self.label == 'local' and not self.__dict__.get('database', None):
182                from meerschaum.config._paths import SQLITE_DB_PATH
183                self.database = str(SQLITE_DB_PATH)
184
185        ### ensure flavor and label are set accordingly
186        if 'flavor' not in self.__dict__:
187            if flavor is None and 'uri' not in self.__dict__:
188                raise Exception(
189                    f"    Missing flavor. Provide flavor as a key for '{self}'."
190                )
191            self.flavor = flavor or self.parse_uri(self.__dict__['uri']).get('flavor', None)
192
193        if self.flavor == 'postgres':
194            self.flavor = 'postgresql'
195
196        self._debug = debug
197        ### Store the PID and thread at initialization
198        ### so we can dispose of the Pool in child processes or threads.
199        import os
200        import threading
201        self._pid = os.getpid()
202        self._thread_ident = threading.current_thread().ident
203        self._sessions = {}
204        self._locks = {'_sessions': threading.RLock(), }
205
206        ### verify the flavor's requirements are met
207        if self.flavor not in self.flavor_configs:
208            error(f"Flavor '{self.flavor}' is not supported by Meerschaum SQLConnector")
209        if not self.__dict__.get('uri'):
210            self.verify_attributes(
211                self.flavor_configs[self.flavor].get('requirements', set()),
212                debug=debug,
213            )
214
215        if wait:
216            from meerschaum.connectors.poll import retry_connect
217            retry_connect(connector=self, debug=debug)
218
219        if connect:
220            if not self.test_connection(debug=debug):
221                warn(f"Failed to connect with connector '{self}'!", stack=False)
222
223    @property
224    def Session(self):
225        if '_Session' not in self.__dict__:
226            if self.engine is None:
227                return None
228
229            from meerschaum.utils.packages import attempt_import
230            sqlalchemy_orm = attempt_import('sqlalchemy.orm', lazy=False)
231            session_factory = sqlalchemy_orm.sessionmaker(self.engine)
232            self._Session = sqlalchemy_orm.scoped_session(session_factory)
233
234        return self._Session
235
236    @property
237    def engine(self):
238        """
239        Return the SQLAlchemy engine connected to the configured database.
240        """
241        import os
242        import threading
243        if '_engine' not in self.__dict__:
244            self._engine, self._engine_str = self.create_engine(include_uri=True)
245
246        same_process = os.getpid() == self._pid
247        same_thread = threading.current_thread().ident == self._thread_ident
248
249        ### handle child processes
250        if not same_process:
251            self._pid = os.getpid()
252            self._thread = threading.current_thread()
253            warn("Different PID detected. Disposing of connections...")
254            self._engine.dispose()
255
256        ### handle different threads
257        if not same_thread:
258            if self.flavor == 'duckdb':
259                warn("Different thread detected.")
260                self._engine.dispose()
261
262        return self._engine
263
264    @property
265    def DATABASE_URL(self) -> str:
266        """
267        Return the URI connection string (alias for `SQLConnector.URI`.
268        """
269        _ = self.engine
270        return str(self._engine_str)
271
272    @property
273    def URI(self) -> str:
274        """
275        Return the URI connection string.
276        """
277        _ = self.engine
278        return str(self._engine_str)
279
280    @property
281    def IS_THREAD_SAFE(self) -> str:
282        """
283        Return whether this connector may be multithreaded.
284        """
285        if self.flavor in ('duckdb', 'oracle'):
286            return False
287        if self.flavor == 'sqlite':
288            return ':memory:' not in self.URI
289        return True
290
291    @property
292    def metadata(self):
293        """
294        Return the metadata bound to this configured schema.
295        """
296        from meerschaum.utils.packages import attempt_import
297        sqlalchemy = attempt_import('sqlalchemy', lazy=False)
298        if '_metadata' not in self.__dict__:
299            self._metadata = sqlalchemy.MetaData(schema=self.schema)
300        return self._metadata
301
302    @property
303    def instance_schema(self):
304        """
305        Return the schema name for Meerschaum tables. 
306        """
307        return self.schema
308
309    @property
310    def internal_schema(self):
311        """
312        Return the schema name for internal tables. 
313        """
314        from meerschaum.config.static import STATIC_CONFIG
315        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
316        schema_name = self.__dict__.get('internal_schema', None) or (
317            STATIC_CONFIG['sql']['internal_schema']
318            if self.flavor not in NO_SCHEMA_FLAVORS
319            else self.schema
320        )
321
322        if '_internal_schema' not in self.__dict__:
323            self._internal_schema = schema_name
324        return self._internal_schema
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    @property
343    def db_version(self) -> Union[str, None]:
344        """
345        Return the database version.
346        """
347        _db_version = self.__dict__.get('_db_version', None)
348        if _db_version is not None:
349            return _db_version
350
351        from meerschaum.utils.sql import get_db_version
352        self._db_version = get_db_version(self)
353        return self._db_version
354
355    @property
356    def schema(self) -> Union[str, None]:
357        """
358        Return the default schema to use.
359        A value of `None` will not prepend a schema.
360        """
361        if 'schema' in self.__dict__:
362            return self.__dict__['schema']
363
364        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
365        if self.flavor in NO_SCHEMA_FLAVORS:
366            self.__dict__['schema'] = None
367            return None
368
369        sqlalchemy = mrsm.attempt_import('sqlalchemy', lazy=False)
370        _schema = sqlalchemy.inspect(self.engine).default_schema_name
371        self.__dict__['schema'] = _schema
372        return _schema
373
374    def __getstate__(self):
375        return self.__dict__
376
377    def __setstate__(self, d):
378        self.__dict__.update(d)
379
380    def __call__(self):
381        return self