meerschaum.connectors

291def make_connector(cls, _is_executor: bool = False):
292    """
293    Register a class as a `Connector`.
294    The `type` will be the lower case of the class name, without the suffix `connector`.
295
296    Parameters
297    ----------
298    instance: bool, default False
299        If `True`, make this connector type an instance connector.
300        This requires implementing the various pipes functions and lots of testing.
301
302    Examples
303    --------
304    >>> import meerschaum as mrsm
305    >>> from meerschaum.connectors import make_connector, Connector
306    >>> 
307    >>> @make_connector
308    >>> class FooConnector(Connector):
309    ...     REQUIRED_ATTRIBUTES: list[str] = ['username', 'password']
310    ... 
311    >>> conn = mrsm.get_connector('foo:bar', username='dog', password='cat')
312    >>> print(conn.username, conn.password)
313    dog cat
314    >>> 
315    """
316    import re
317    suffix_regex = (
318        r'connector$'
319        if not _is_executor
320        else r'executor$'
321    )
322    typ = re.sub(suffix_regex, '', cls.__name__.lower())
323    with _locks['types']:
324        types[typ] = cls
325    with _locks['custom_types']:
326        custom_types.add(typ)
327    with _locks['connectors']:
328        if typ not in connectors:
329            connectors[typ] = {}
330    if getattr(cls, 'IS_INSTANCE', False):
331        with _locks['instance_types']:
332            if typ not in instance_types:
333                instance_types.append(typ)
334
335    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            if uri.startswith('postgis://'):
156                uri = uri.replace('postgis://', 'postgresql+psycopg://', 1)
157                flavor = 'postgis'
158            kw['uri'] = uri
159            from_uri_params = self.from_uri(kw['uri'], as_dict=True)
160            label = label or from_uri_params.get('label', None)
161            _ = from_uri_params.pop('label', None)
162
163            ### Sometimes the flavor may be provided with a URI.
164            kw.update(from_uri_params)
165            if flavor:
166                kw['flavor'] = flavor
167
168        ### set __dict__ in base class
169        super().__init__(
170            'sql',
171            label = label or self.__dict__.get('label', None),
172            **kw
173        )
174
175        if self.__dict__.get('flavor', None) == 'sqlite':
176            self._reset_attributes()
177            self._set_attributes(
178                'sql',
179                label = label,
180                inherit_default = False,
181                **kw
182            )
183            ### For backwards compatability reasons, set the path for sql:local if its missing.
184            if self.label == 'local' and not self.__dict__.get('database', None):
185                from meerschaum.config._paths import SQLITE_DB_PATH
186                self.database = str(SQLITE_DB_PATH)
187
188        ### ensure flavor and label are set accordingly
189        if 'flavor' not in self.__dict__:
190            if flavor is None and 'uri' not in self.__dict__:
191                raise ValueError(
192                    f"    Missing flavor. Provide flavor as a key for '{self}'."
193                )
194            self.flavor = flavor or self.parse_uri(self.__dict__['uri']).get('flavor', None)
195
196        if self.flavor == 'postgres':
197            self.flavor = 'postgresql'
198
199        self._debug = debug
200        ### Store the PID and thread at initialization
201        ### so we can dispose of the Pool in child processes or threads.
202        import os
203        import threading
204        self._pid = os.getpid()
205        self._thread_ident = threading.current_thread().ident
206        self._sessions = {}
207        self._locks = {'_sessions': threading.RLock(), }
208
209        ### verify the flavor's requirements are met
210        if self.flavor not in self.flavor_configs:
211            error(f"Flavor '{self.flavor}' is not supported by Meerschaum SQLConnector")
212        if not self.__dict__.get('uri'):
213            self.verify_attributes(
214                self.flavor_configs[self.flavor].get('requirements', set()),
215                debug=debug,
216            )
217
218        if wait:
219            from meerschaum.connectors.poll import retry_connect
220            retry_connect(connector=self, debug=debug)
221
222        if connect:
223            if not self.test_connection(debug=debug):
224                warn(f"Failed to connect with connector '{self}'!", stack=False)
225
226    @property
227    def Session(self):
228        if '_Session' not in self.__dict__:
229            if self.engine is None:
230                return None
231
232            from meerschaum.utils.packages import attempt_import
233            sqlalchemy_orm = attempt_import('sqlalchemy.orm', lazy=False)
234            session_factory = sqlalchemy_orm.sessionmaker(self.engine)
235            self._Session = sqlalchemy_orm.scoped_session(session_factory)
236
237        return self._Session
238
239    @property
240    def engine(self):
241        """
242        Return the SQLAlchemy engine connected to the configured database.
243        """
244        import os
245        import threading
246        if '_engine' not in self.__dict__:
247            self._engine, self._engine_str = self.create_engine(include_uri=True)
248
249        same_process = os.getpid() == self._pid
250        same_thread = threading.current_thread().ident == self._thread_ident
251
252        ### handle child processes
253        if not same_process:
254            self._pid = os.getpid()
255            self._thread = threading.current_thread()
256            warn("Different PID detected. Disposing of connections...")
257            self._engine.dispose()
258
259        ### handle different threads
260        if not same_thread:
261            if self.flavor == 'duckdb':
262                warn("Different thread detected.")
263                self._engine.dispose()
264
265        return self._engine
266
267    @property
268    def DATABASE_URL(self) -> str:
269        """
270        Return the URI connection string (alias for `SQLConnector.URI`.
271        """
272        _ = self.engine
273        return str(self._engine_str)
274
275    @property
276    def URI(self) -> str:
277        """
278        Return the URI connection string.
279        """
280        _ = self.engine
281        return str(self._engine_str)
282
283    @property
284    def IS_THREAD_SAFE(self) -> str:
285        """
286        Return whether this connector may be multithreaded.
287        """
288        if self.flavor in ('duckdb', 'oracle'):
289            return False
290        if self.flavor == 'sqlite':
291            return ':memory:' not in self.URI
292        return True
293
294    @property
295    def metadata(self):
296        """
297        Return the metadata bound to this configured schema.
298        """
299        from meerschaum.utils.packages import attempt_import
300        sqlalchemy = attempt_import('sqlalchemy', lazy=False)
301        if '_metadata' not in self.__dict__:
302            self._metadata = sqlalchemy.MetaData(schema=self.schema)
303        return self._metadata
304
305    @property
306    def instance_schema(self):
307        """
308        Return the schema name for Meerschaum tables. 
309        """
310        return self.schema
311
312    @property
313    def internal_schema(self):
314        """
315        Return the schema name for internal tables. 
316        """
317        from meerschaum.config.static import STATIC_CONFIG
318        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
319        schema_name = self.__dict__.get('internal_schema', None) or (
320            STATIC_CONFIG['sql']['internal_schema']
321            if self.flavor not in NO_SCHEMA_FLAVORS
322            else self.schema
323        )
324
325        if '_internal_schema' not in self.__dict__:
326            self._internal_schema = schema_name
327        return self._internal_schema
328
329    @property
330    def db(self) -> Optional[databases.Database]:
331        from meerschaum.utils.packages import attempt_import
332        databases = attempt_import('databases', lazy=False, install=True)
333        url = self.DATABASE_URL
334        if 'mysql' in url:
335            url = url.replace('+pymysql', '')
336        if '_db' not in self.__dict__:
337            try:
338                self._db = databases.Database(url)
339            except KeyError:
340                ### Likely encountered an unsupported flavor.
341                from meerschaum.utils.warnings import warn
342                self._db = None
343        return self._db
344
345    @property
346    def db_version(self) -> Union[str, None]:
347        """
348        Return the database version.
349        """
350        _db_version = self.__dict__.get('_db_version', None)
351        if _db_version is not None:
352            return _db_version
353
354        from meerschaum.utils.sql import get_db_version
355        self._db_version = get_db_version(self)
356        return self._db_version
357
358    @property
359    def schema(self) -> Union[str, None]:
360        """
361        Return the default schema to use.
362        A value of `None` will not prepend a schema.
363        """
364        if 'schema' in self.__dict__:
365            return self.__dict__['schema']
366
367        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
368        if self.flavor in NO_SCHEMA_FLAVORS:
369            self.__dict__['schema'] = None
370            return None
371
372        sqlalchemy = mrsm.attempt_import('sqlalchemy', lazy=False)
373        _schema = sqlalchemy.inspect(self.engine).default_schema_name
374        self.__dict__['schema'] = _schema
375        return _schema
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