meerschaum.connectors

  1#! /usr/bin/env python
  2# -*- coding: utf-8 -*-
  3# vim:fenc=utf-8
  4
  5"""
  6Create connectors with `meerschaum.connectors.get_connector()`.
  7For ease of use, you can also import from the root `meerschaum` module:
  8```
  9>>> from meerschaum import get_connector
 10>>> conn = get_connector()
 11```
 12"""
 13
 14from __future__ import annotations
 15
 16import meerschaum as mrsm
 17from meerschaum.utils.typing import Any, Union, List, Dict, Optional
 18from meerschaum.utils.threading import RLock
 19from meerschaum.utils.warnings import warn
 20
 21from meerschaum.connectors._Connector import Connector, InvalidAttributesError
 22from meerschaum.connectors.instance._InstanceConnector import InstanceConnector
 23from meerschaum.connectors.sql._SQLConnector import SQLConnector
 24from meerschaum.connectors.api._APIConnector import APIConnector
 25
 26__all__ = (
 27    "make_connector",
 28    "Connector",
 29    "InstanceConnector",
 30    "SQLConnector",
 31    "APIConnector",
 32    "get_connector",
 33    "is_connected",
 34    "allow_protected",
 35    "poll",
 36    "api",
 37    "sql",
 38    "valkey",
 39    "parse",
 40)
 41
 42### store connectors partitioned by
 43### type, label for reuse
 44connectors: Dict[str, Dict[str, Connector]] = {
 45    'api'    : {},
 46    'sql'    : {},
 47    'plugin' : {},
 48    'valkey' : {},
 49}
 50instance_types: List[str] = ['sql', 'api']
 51_locks: Dict[str, RLock] = {
 52    'connectors'               : RLock(),
 53    'types'                    : RLock(),
 54    'custom_types'             : RLock(),
 55    'plugins_types'            : RLock(),
 56    '_loaded_plugin_connectors': RLock(),
 57    'instance_types'           : RLock(),
 58    'native_types'             : RLock(),
 59}
 60
 61### Fill this with objects only when connectors are first referenced.
 62types: Dict[str, Any] = {}
 63custom_types: set = set()
 64plugins_types: Dict[str, List[str]] = {}
 65_known_custom_types: set = set()
 66_loaded_plugin_connectors: bool = False
 67
 68### Connector types declared by native (Rust) plugins: `{type: info}` from the core, read on the
 69### first lookup of a type Python does not know, and the plugin of each type registered here.
 70_native_types: Optional[Dict[str, Dict[str, Any]]] = None
 71_native_types_refresh: bool = False
 72_native_type_plugins: Dict[str, str] = {}
 73
 74
 75def get_connector(
 76    type: str = None,
 77    label: str = None,
 78    refresh: bool = False,
 79    debug: bool = False,
 80    _load_plugins: bool = True,
 81    _warn_undefined: bool = True,
 82    **kw: Any
 83) -> Connector:
 84    """
 85    Return existing connector or create new connection and store for reuse.
 86    
 87    You can create new connectors if enough parameters are provided for the given type and flavor.
 88
 89    Parameters
 90    ----------
 91    type: Optional[str], default None
 92        Connector type (sql, api, etc.).
 93        Defaults to the type of the configured `instance_connector`.
 94
 95    label: Optional[str], default None
 96        Connector label (e.g. main). Defaults to `'main'`.
 97
 98    refresh: bool, default False
 99        Refresh the Connector instance / construct new object. Defaults to `False`.
100
101    kw: Any
102        Other arguments to pass to the Connector constructor.
103        If the Connector has already been constructed and new arguments are provided,
104        `refresh` is set to `True` and the old Connector is replaced.
105
106    Returns
107    -------
108    A new Meerschaum connector (e.g. `meerschaum.connectors.api.APIConnector`,
109    `meerschaum.connectors.sql.SQLConnector`).
110    
111    Examples
112    --------
113    The following parameters would create a new
114    `meerschaum.connectors.sql.SQLConnector` that isn't in the configuration file.
115
116    ```
117    >>> conn = get_connector(
118    ...     type = 'sql',
119    ...     label = 'newlabel',
120    ...     flavor = 'sqlite',
121    ...     database = '/file/path/to/database.db'
122    ... )
123    >>>
124    ```
125
126    """
127    from meerschaum.connectors.parse import parse_instance_keys
128    from meerschaum.config import get_config
129    from meerschaum._internal.static import STATIC_CONFIG
130    from meerschaum.utils.warnings import warn
131    global _loaded_plugin_connectors
132    if isinstance(type, str) and not label and ':' in type:
133        type, label = type.split(':', maxsplit=1)
134
135    if _load_plugins:
136        with _locks['_loaded_plugin_connectors']:
137            if not _loaded_plugin_connectors:
138                load_plugin_connectors()
139                _load_builtin_custom_connectors()
140                _loaded_plugin_connectors = True
141
142    if type is None and label is None:
143        default_instance_keys = get_config('meerschaum', 'instance', patch=True)
144        ### recursive call to get_connector
145        return parse_instance_keys(default_instance_keys)
146
147    ### NOTE: the default instance connector may not be main.
148    ### Only fall back to 'main' if the type is provided by the label is omitted.
149    label = label if label is not None else STATIC_CONFIG['connectors']['default_label']
150
151    ### A type declared by a native (Rust) plugin is registered on first use.
152    if type not in connectors and _load_plugins and isinstance(type, str):
153        _load_native_connector_type(type)
154
155    ### type might actually be a label. Check if so and raise a warning.
156    if type not in connectors:
157        possibilities, poss_msg = [], ""
158        for _type in get_config('meerschaum', 'connectors'):
159            if type in get_config('meerschaum', 'connectors', _type):
160                possibilities.append(f"{_type}:{type}")
161        if len(possibilities) > 0:
162            poss_msg = " Did you mean"
163            for poss in possibilities[:-1]:
164                poss_msg += f" '{poss}',"
165            if poss_msg.endswith(','):
166                poss_msg = poss_msg[:-1]
167            if len(possibilities) > 1:
168                poss_msg += " or"
169            poss_msg += f" '{possibilities[-1]}'?"
170
171        warn(f"Cannot create Connector of type '{type}'." + poss_msg, stack=False)
172        return None
173
174    if 'sql' not in types:
175        from meerschaum.connectors.plugin import PluginConnector
176        from meerschaum.connectors.valkey import ValkeyConnector
177        with _locks['types']:
178            types.update({
179                'api': APIConnector,
180                'sql': SQLConnector,
181                'plugin': PluginConnector,
182                'valkey': ValkeyConnector,
183            })
184
185    ### determine if we need to call the constructor
186    if not refresh:
187        ### see if any user-supplied arguments differ from the existing instance
188        if label in connectors[type]:
189            warning_message = None
190            for attribute, value in kw.items():
191                if attribute not in connectors[type][label].meta:
192                    import inspect
193                    cls = connectors[type][label].__class__
194                    cls_init_signature = inspect.signature(cls)
195                    cls_init_params = cls_init_signature.parameters
196                    if attribute not in cls_init_params:
197                        warning_message = (
198                            f"Received new attribute '{attribute}' not present in connector " +
199                            f"{connectors[type][label]}.\n"
200                        )
201                elif connectors[type][label].__dict__[attribute] != value:
202                    warning_message = (
203                        f"Mismatched values for attribute '{attribute}' in connector "
204                        + f"'{connectors[type][label]}'.\n" +
205                        f"  - Keyword value: '{value}'\n" +
206                        f"  - Existing value: '{connectors[type][label].__dict__[attribute]}'\n"
207                    )
208            if warning_message is not None:
209                warning_message += (
210                    "\nSetting `refresh` to True and recreating connector with type:"
211                    + f" '{type}' and label '{label}'."
212                )
213                refresh = True
214                warn(warning_message)
215        else: ### connector doesn't yet exist
216            refresh = True
217
218    ### only create an object if refresh is True
219    ### (can be manually specified, otherwise determined above)
220    if refresh:
221        ### A connector defined nowhere: say where connectors come from (not "Missing flavor").
222        undefined = _undefined_connector_message(type, label, kw)
223        if undefined is not None:
224            if _warn_undefined:
225                warn(undefined, stack=False)
226            return None
227        with _locks['connectors']:
228            try:
229                ### will raise an error if configuration is incorrect / missing
230                conn = types[type](label=label, **kw)
231                connectors[type][label] = conn
232            except InvalidAttributesError as ie:
233                warn(
234                    f"Incorrect attributes for connector '{type}:{label}'.\n"
235                    + str(ie),
236                    stack = False,
237                )
238                conn = None
239            except Exception as e:
240                from meerschaum.utils.formatting import print_exception
241                print_exception()
242                warn(
243                    f"Exception when creating connector '{type}:{label}'.\n" + str(e),
244                    stack = False,
245                )
246                conn = None
247        if conn is None:
248            return None
249
250    return connectors[type][label]
251
252
253def _undefined_connector_message(
254    type: Optional[str],
255    label: Optional[str],
256    attributes: Optional[Dict[str, Any]] = None,
257) -> Optional[str]:
258    """
259    The message for a `sql:`, `api:` or `valkey:` connector defined nowhere: no attributes were
260    passed and the config has no section for it (which includes the `MRSM_<TYPE>_<LABEL>`
261    environment connectors). `None` when the connector is defined. The text is the CLI's
262    (`undefined_connector_message` in `crates/meerschaum-cli/src/native/connectors.rs`).
263    """
264    from meerschaum.config import get_config
265    if type not in ('sql', 'api', 'valkey') or not label or attributes:
266        return None
267    if get_config('meerschaum', 'connectors', type, label, warn=False):
268        return None
269    return (
270        f"Connector '{type}:{label}' is not defined. Add it to the config "
271        f"(`mrsm bootstrap connectors`), set `MRSM_{type.upper()}_{label.upper()}`, "
272        "or install the plugin that provides it."
273    )
274
275
276class allow_protected:
277    """
278    Allow writes to protected connectors (`protected: true`) inside a `with` block.
279
280    Every mutation the core makes through a protected connector (sync, clear, drop, delete,
281    register, edit, `verify schema`, `SQLConnector.exec()` / `to_sql()`) is refused unless writes
282    are allowed: by `--allow-protected` on the command line, by `MRSM_ALLOW_PROTECTED=1` in the
283    environment, or inside this context manager. The allowance is process-wide while the block
284    runs (other threads see it too); child processes need `MRSM_ALLOW_PROTECTED=1`.
285
286    Examples
287    --------
288    >>> import meerschaum as mrsm
289    >>> from meerschaum.connectors import allow_protected
290    >>> pipe = mrsm.Pipe('plugin:noaa', 'weather', instance='sql:prod')
291    >>> with allow_protected():
292    ...     success, msg = pipe.sync(df)
293    """
294
295    def __enter__(self) -> 'allow_protected':
296        from meerschaum import _core
297        _core.allow_protected_enter()
298        return self
299
300    def __exit__(self, *exc: Any) -> None:
301        from meerschaum import _core
302        _core.allow_protected_exit()
303
304
305def is_protected_value(value: Any) -> bool:
306    """
307    Whether a connector attribute value marks the connector protected
308    (`True`, a non-zero number, or a string like `'true'`, `'1'`, `'yes'`).
309    """
310    if isinstance(value, bool):
311        return value
312    if isinstance(value, (int, float)):
313        return value != 0
314    if isinstance(value, str):
315        return value.strip().lower() in ('1', 'true', 'yes', 'on', 'y', 't')
316    return False
317
318
319def uri_marks_protected(uri: Any) -> bool:
320    """
321    Whether a connector URI carries `?protected=true` (any truthy value), which marks the
322    connector protected, e.g. `MRSM_VALKEY_PROD='valkey://host:6379/0?protected=true'`.
323    """
324    if not isinstance(uri, str) or '?' not in uri:
325        return False
326    from meerschaum import _core
327    return _core.uri_marks_protected(uri)
328
329
330def protected_refusal(connector: Any, op: str, pipe: Any = None) -> Optional[str]:
331    """
332    The refusal message for a write `op` through `connector` when it is protected and writes
333    are not allowed, else `None`. Meerschaum's own internal pipes (connector keys `mrsm`) are
334    never refused.
335    """
336    if pipe is not None and str(getattr(pipe, 'connector_keys', '')) == 'mrsm':
337        return None
338    if not is_protected_value(getattr(connector, 'protected', False)):
339        return None
340    from meerschaum import _core
341    if _core.protected_allowed():
342        return None
343    what = f"{op} {pipe}" if pipe is not None else op
344    return _core.protected_refusal_message(what, str(connector))
345
346
347def is_connected(keys: str, **kw) -> bool:
348    """
349    Check if the connector keys correspond to an active connection.
350    If the connector has not been created, it will immediately return `False`.
351    If the connector exists but cannot communicate with the source, return `False`.
352    
353    **NOTE:** Only works with instance connectors (`SQLConnectors` and `APIConnectors`).
354    Keyword arguments are passed to `meerschaum.connectors.poll.retry_connect`.
355
356    Parameters
357    ----------
358    keys:
359        The keys to the connector (e.g. `'sql:main'`).
360        
361    Returns
362    -------
363    A `bool` corresponding to whether a successful connection may be made.
364
365    """
366    import warnings
367    if ':' not in keys:
368        warn(f"Invalid connector keys '{keys}'")
369
370    try:
371        typ, label = keys.split(':')
372    except Exception:
373        return False
374    if typ not in instance_types:
375        return False
376    if label not in connectors.get(typ, {}):
377        return False
378
379    from meerschaum.connectors.parse import parse_instance_keys
380    conn = parse_instance_keys(keys)
381    try:
382        with warnings.catch_warnings():
383            warnings.filterwarnings('ignore')
384            return conn.test_connection(**kw)
385    except Exception:
386        return False
387
388
389def make_connector(cls, _is_executor: bool = False):
390    """
391    Register a class as a `Connector`.
392    The `type` will be the lower case of the class name, without the suffix `connector`.
393
394    Parameters
395    ----------
396    instance: bool, default False
397        If `True`, make this connector type an instance connector.
398        This requires implementing the various pipes functions and lots of testing.
399
400    Examples
401    --------
402    >>> import meerschaum as mrsm
403    >>> from meerschaum.connectors import make_connector, Connector
404    >>> 
405    >>> @make_connector
406    >>> class FooConnector(Connector):
407    ...     REQUIRED_ATTRIBUTES: list[str] = ['username', 'password']
408    ... 
409    >>> conn = mrsm.get_connector('foo:bar', username='dog', password='cat')
410    >>> print(conn.username, conn.password)
411    dog cat
412    >>> 
413    """
414    import re
415    from meerschaum.plugins import _get_parent_plugin
416    suffix_regex = (
417        r'connector$'
418        if not _is_executor
419        else r'executor$'
420    )
421    plugin_name = _get_parent_plugin(2)
422    typ = re.sub(suffix_regex, '', cls.__name__.lower())
423    with _locks['types']:
424        types[typ] = cls
425    with _locks['custom_types']:
426        custom_types.add(typ)
427    if plugin_name:
428        with _locks['plugins_types']:
429            if plugin_name not in plugins_types:
430                plugins_types[plugin_name] = []
431            plugins_types[plugin_name].append(typ)
432    with _locks['connectors']:
433        if typ not in connectors:
434            connectors[typ] = {}
435    if getattr(cls, 'IS_INSTANCE', False):
436        with _locks['instance_types']:
437            if typ not in instance_types:
438                instance_types.append(typ)
439
440    return cls
441
442
443def load_plugin_connectors():
444    """
445    If a plugin makes use of the `make_connector` decorator,
446    load its module.
447    """
448    import os
449    from meerschaum.plugins import get_plugins, import_plugins
450    to_import = []
451    for plugin in get_plugins():
452        if plugin is None:
453            continue
454
455        plugin_file = plugin.__file__
456        if not plugin_file:
457            continue
458
459        files_to_scan = []
460        plugin_dir = os.path.dirname(plugin_file)
461        if os.path.basename(plugin_file) == '__init__.py':
462            for root, _, files in os.walk(plugin_dir):
463                for fname in files:
464                    if fname.endswith('.py'):
465                        files_to_scan.append(os.path.join(root, fname))
466        else:
467            files_to_scan.append(plugin_file)
468
469        found = False
470        for fpath in files_to_scan:
471            try:
472                with open(fpath, encoding='utf-8') as f:
473                    text = f.read()
474            except (OSError, UnicodeDecodeError):
475                continue
476            if 'make_connector' in text or 'Connector' in text:
477                found = True
478                break
479
480        if found:
481            to_import.append(plugin.name)
482
483    if not to_import:
484        return
485
486    import_plugins(*to_import)
487
488
489def unload_plugin_connectors(
490    plugin_names: Optional[List[str]] = None,
491    debug: bool = False,
492) -> None:
493    """
494    Unload custom connectors added by plugins.
495    """
496    from meerschaum.plugins import get_plugins_names
497    global custom_types, _known_custom_types, types, plugins_types, connectors, _loaded_plugin_connectors
498
499    plugin_names = plugin_names or get_plugins_names()
500
501    for plugin_name in plugin_names:
502        plugin_types = plugins_types.get(plugin_name, [])
503        for typ in plugin_types:
504            _ = types.pop(typ, None)
505            _ = connectors.pop(typ, None)
506            if typ in instance_types:
507                instance_types.remove(typ)
508
509    custom_types.clear()
510    custom_types.update(_known_custom_types)
511    _loaded_plugin_connectors = False
512    _unload_native_connector_types()
513
514
515def get_connector_plugin(
516    connector: Connector,
517) -> Union[str, None, mrsm.Plugin]:
518    """
519    Determine the plugin for a connector.
520    This is useful for handling virtual environments for custom instance connectors.
521
522    Parameters
523    ----------
524    connector: Connector
525        The connector which may require a virtual environment.
526
527    Returns
528    -------
529    A Plugin, 'mrsm', or None.
530    """
531    if not hasattr(connector, 'type'):
532        return None
533    ### A native plugin has no virtual environment.
534    if getattr(connector, 'IS_NATIVE', False):
535        return None
536    plugin_name = (
537        connector.__module__.replace('plugins.', '').split('.')[0]
538        if connector.type in custom_types else (
539            connector.label
540            if connector.type == 'plugin'
541            else 'mrsm'
542        )
543    )
544    plugin = mrsm.Plugin(plugin_name)
545    return plugin if plugin.is_installed() else None
546
547
548def _load_builtin_custom_connectors():
549    """
550    Import custom connectors decorated with `@make_connector` or `@make_executor`.
551    """
552    import meerschaum.jobs.systemd
553    import meerschaum.connectors.valkey
554    _known_custom_types.add('valkey')
555    _known_custom_types.add('systemd')
556
557
558def get_native_connector_types(refresh: bool = False) -> Dict[str, Dict[str, Any]]:
559    """
560    Return the connector types declared by native (Rust) plugins in the plugins directories,
561    as `{type: {'plugin', 'is_instance', 'is_streaming', 'is_thread_safe', 'methods'}}`.
562    Read once per process (like Python plugins' connector types) until
563    `unload_plugin_connectors()` or `refresh=True`.
564    """
565    global _native_types, _native_types_refresh
566    with _locks['native_types']:
567        if _native_types is not None and not refresh:
568            return _native_types
569        from meerschaum import _core
570        try:
571            _native_types = {
572                str(typ): dict(info)
573                for typ, info in _core.native_connector_types(
574                    refresh or _native_types_refresh
575                ).items()
576            }
577        except Exception as e:
578            from meerschaum.utils.debug import dprint
579            dprint(f"Could not read the native plugins' connector types: {e}")
580            _native_types = {}
581        _native_types_refresh = False
582        return _native_types
583
584
585def _load_native_connector_type(typ: str) -> bool:
586    """
587    Register the native connector type `typ` (a `NativeConnector` class) if a native plugin
588    declares it. Return whether `typ` is now a known type.
589    """
590    info = get_native_connector_types().get(typ, None)
591    if not info:
592        return False
593    from meerschaum.connectors.native import make_native_connector_class
594    cls = make_native_connector_class(typ, info)
595    plugin_name = str(info.get('plugin', ''))
596    with _locks['types']:
597        if typ in types:
598            return True
599        types[typ] = cls
600    with _locks['custom_types']:
601        custom_types.add(typ)
602    with _locks['plugins_types']:
603        plugins_types.setdefault(plugin_name, []).append(typ)
604    _native_type_plugins[typ] = plugin_name
605    with _locks['connectors']:
606        if typ not in connectors:
607            connectors[typ] = {}
608    if getattr(cls, 'IS_INSTANCE', False):
609        with _locks['instance_types']:
610            if typ not in instance_types:
611                instance_types.append(typ)
612    return True
613
614
615def _unload_native_connector_types() -> None:
616    """
617    Forget the native connector types (re-read from the core on the next unknown type).
618    """
619    global _native_types, _native_types_refresh
620    with _locks['native_types']:
621        for typ, plugin_name in list(_native_type_plugins.items()):
622            _ = types.pop(typ, None)
623            _ = connectors.pop(typ, None)
624            custom_types.discard(typ)
625            if typ in instance_types:
626                instance_types.remove(typ)
627            plugin_types = plugins_types.get(plugin_name, [])
628            if typ in plugin_types:
629                plugin_types.remove(typ)
630            _ = _native_type_plugins.pop(typ, None)
631        _native_types = None
632        _native_types_refresh = True