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    "poll",
 35    "api",
 36    "sql",
 37    "valkey",
 38    "parse",
 39)
 40
 41### store connectors partitioned by
 42### type, label for reuse
 43connectors: Dict[str, Dict[str, Connector]] = {
 44    'api'    : {},
 45    'sql'    : {},
 46    'plugin' : {},
 47    'valkey' : {},
 48}
 49instance_types: List[str] = ['sql', 'api']
 50_locks: Dict[str, RLock] = {
 51    'connectors'               : RLock(),
 52    'types'                    : RLock(),
 53    'custom_types'             : RLock(),
 54    'plugins_types'            : RLock(),
 55    '_loaded_plugin_connectors': RLock(),
 56    'instance_types'           : RLock(),
 57}
 58
 59### Fill this with objects only when connectors are first referenced.
 60types: Dict[str, Any] = {}
 61custom_types: set = set()
 62plugins_types: Dict[str, List[str]] = {}
 63_known_custom_types: set = set()
 64_loaded_plugin_connectors: bool = False
 65
 66
 67def get_connector(
 68    type: str = None,
 69    label: str = None,
 70    refresh: bool = False,
 71    debug: bool = False,
 72    _load_plugins: bool = True,
 73    **kw: Any
 74) -> Connector:
 75    """
 76    Return existing connector or create new connection and store for reuse.
 77    
 78    You can create new connectors if enough parameters are provided for the given type and flavor.
 79
 80    Parameters
 81    ----------
 82    type: Optional[str], default None
 83        Connector type (sql, api, etc.).
 84        Defaults to the type of the configured `instance_connector`.
 85
 86    label: Optional[str], default None
 87        Connector label (e.g. main). Defaults to `'main'`.
 88
 89    refresh: bool, default False
 90        Refresh the Connector instance / construct new object. Defaults to `False`.
 91
 92    kw: Any
 93        Other arguments to pass to the Connector constructor.
 94        If the Connector has already been constructed and new arguments are provided,
 95        `refresh` is set to `True` and the old Connector is replaced.
 96
 97    Returns
 98    -------
 99    A new Meerschaum connector (e.g. `meerschaum.connectors.api.APIConnector`,
100    `meerschaum.connectors.sql.SQLConnector`).
101    
102    Examples
103    --------
104    The following parameters would create a new
105    `meerschaum.connectors.sql.SQLConnector` that isn't in the configuration file.
106
107    ```
108    >>> conn = get_connector(
109    ...     type = 'sql',
110    ...     label = 'newlabel',
111    ...     flavor = 'sqlite',
112    ...     database = '/file/path/to/database.db'
113    ... )
114    >>>
115    ```
116
117    """
118    from meerschaum.connectors.parse import parse_instance_keys
119    from meerschaum.config import get_config
120    from meerschaum._internal.static import STATIC_CONFIG
121    from meerschaum.utils.warnings import warn
122    global _loaded_plugin_connectors
123    if isinstance(type, str) and not label and ':' in type:
124        type, label = type.split(':', maxsplit=1)
125
126    if _load_plugins:
127        with _locks['_loaded_plugin_connectors']:
128            if not _loaded_plugin_connectors:
129                load_plugin_connectors()
130                _load_builtin_custom_connectors()
131                _loaded_plugin_connectors = True
132
133    if type is None and label is None:
134        default_instance_keys = get_config('meerschaum', 'instance', patch=True)
135        ### recursive call to get_connector
136        return parse_instance_keys(default_instance_keys)
137
138    ### NOTE: the default instance connector may not be main.
139    ### Only fall back to 'main' if the type is provided by the label is omitted.
140    label = label if label is not None else STATIC_CONFIG['connectors']['default_label']
141
142    ### type might actually be a label. Check if so and raise a warning.
143    if type not in connectors:
144        possibilities, poss_msg = [], ""
145        for _type in get_config('meerschaum', 'connectors'):
146            if type in get_config('meerschaum', 'connectors', _type):
147                possibilities.append(f"{_type}:{type}")
148        if len(possibilities) > 0:
149            poss_msg = " Did you mean"
150            for poss in possibilities[:-1]:
151                poss_msg += f" '{poss}',"
152            if poss_msg.endswith(','):
153                poss_msg = poss_msg[:-1]
154            if len(possibilities) > 1:
155                poss_msg += " or"
156            poss_msg += f" '{possibilities[-1]}'?"
157
158        warn(f"Cannot create Connector of type '{type}'." + poss_msg, stack=False)
159        return None
160
161    if 'sql' not in types:
162        from meerschaum.connectors.plugin import PluginConnector
163        from meerschaum.connectors.valkey import ValkeyConnector
164        with _locks['types']:
165            types.update({
166                'api': APIConnector,
167                'sql': SQLConnector,
168                'plugin': PluginConnector,
169                'valkey': ValkeyConnector,
170            })
171
172    ### determine if we need to call the constructor
173    if not refresh:
174        ### see if any user-supplied arguments differ from the existing instance
175        if label in connectors[type]:
176            warning_message = None
177            for attribute, value in kw.items():
178                if attribute not in connectors[type][label].meta:
179                    import inspect
180                    cls = connectors[type][label].__class__
181                    cls_init_signature = inspect.signature(cls)
182                    cls_init_params = cls_init_signature.parameters
183                    if attribute not in cls_init_params:
184                        warning_message = (
185                            f"Received new attribute '{attribute}' not present in connector " +
186                            f"{connectors[type][label]}.\n"
187                        )
188                elif connectors[type][label].__dict__[attribute] != value:
189                    warning_message = (
190                        f"Mismatched values for attribute '{attribute}' in connector "
191                        + f"'{connectors[type][label]}'.\n" +
192                        f"  - Keyword value: '{value}'\n" +
193                        f"  - Existing value: '{connectors[type][label].__dict__[attribute]}'\n"
194                    )
195            if warning_message is not None:
196                warning_message += (
197                    "\nSetting `refresh` to True and recreating connector with type:"
198                    + f" '{type}' and label '{label}'."
199                )
200                refresh = True
201                warn(warning_message)
202        else: ### connector doesn't yet exist
203            refresh = True
204
205    ### only create an object if refresh is True
206    ### (can be manually specified, otherwise determined above)
207    if refresh:
208        with _locks['connectors']:
209            try:
210                ### will raise an error if configuration is incorrect / missing
211                conn = types[type](label=label, **kw)
212                connectors[type][label] = conn
213            except InvalidAttributesError as ie:
214                warn(
215                    f"Incorrect attributes for connector '{type}:{label}'.\n"
216                    + str(ie),
217                    stack = False,
218                )
219                conn = None
220            except Exception as e:
221                from meerschaum.utils.formatting import get_console
222                console = get_console()
223                if console:
224                    console.print_exception()
225                warn(
226                    f"Exception when creating connector '{type}:{label}'.\n" + str(e),
227                    stack = False,
228                )
229                conn = None
230        if conn is None:
231            return None
232
233    return connectors[type][label]
234
235
236def is_connected(keys: str, **kw) -> bool:
237    """
238    Check if the connector keys correspond to an active connection.
239    If the connector has not been created, it will immediately return `False`.
240    If the connector exists but cannot communicate with the source, return `False`.
241    
242    **NOTE:** Only works with instance connectors (`SQLConnectors` and `APIConnectors`).
243    Keyword arguments are passed to `meerschaum.connectors.poll.retry_connect`.
244
245    Parameters
246    ----------
247    keys:
248        The keys to the connector (e.g. `'sql:main'`).
249        
250    Returns
251    -------
252    A `bool` corresponding to whether a successful connection may be made.
253
254    """
255    import warnings
256    if ':' not in keys:
257        warn(f"Invalid connector keys '{keys}'")
258
259    try:
260        typ, label = keys.split(':')
261    except Exception:
262        return False
263    if typ not in instance_types:
264        return False
265    if label not in connectors.get(typ, {}):
266        return False
267
268    from meerschaum.connectors.parse import parse_instance_keys
269    conn = parse_instance_keys(keys)
270    try:
271        with warnings.catch_warnings():
272            warnings.filterwarnings('ignore')
273            return conn.test_connection(**kw)
274    except Exception:
275        return False
276
277
278def make_connector(cls, _is_executor: bool = False):
279    """
280    Register a class as a `Connector`.
281    The `type` will be the lower case of the class name, without the suffix `connector`.
282
283    Parameters
284    ----------
285    instance: bool, default False
286        If `True`, make this connector type an instance connector.
287        This requires implementing the various pipes functions and lots of testing.
288
289    Examples
290    --------
291    >>> import meerschaum as mrsm
292    >>> from meerschaum.connectors import make_connector, Connector
293    >>> 
294    >>> @make_connector
295    >>> class FooConnector(Connector):
296    ...     REQUIRED_ATTRIBUTES: list[str] = ['username', 'password']
297    ... 
298    >>> conn = mrsm.get_connector('foo:bar', username='dog', password='cat')
299    >>> print(conn.username, conn.password)
300    dog cat
301    >>> 
302    """
303    import re
304    from meerschaum.plugins import _get_parent_plugin
305    suffix_regex = (
306        r'connector$'
307        if not _is_executor
308        else r'executor$'
309    )
310    plugin_name = _get_parent_plugin(2)
311    typ = re.sub(suffix_regex, '', cls.__name__.lower())
312    with _locks['types']:
313        types[typ] = cls
314    with _locks['custom_types']:
315        custom_types.add(typ)
316    if plugin_name:
317        with _locks['plugins_types']:
318            if plugin_name not in plugins_types:
319                plugins_types[plugin_name] = []
320            plugins_types[plugin_name].append(typ)
321    with _locks['connectors']:
322        if typ not in connectors:
323            connectors[typ] = {}
324    if getattr(cls, 'IS_INSTANCE', False):
325        with _locks['instance_types']:
326            if typ not in instance_types:
327                instance_types.append(typ)
328
329    return cls
330
331
332def load_plugin_connectors():
333    """
334    If a plugin makes use of the `make_connector` decorator,
335    load its module.
336    """
337    import os
338    from meerschaum.plugins import get_plugins, import_plugins
339    to_import = []
340    for plugin in get_plugins():
341        if plugin is None:
342            continue
343
344        plugin_file = plugin.__file__
345        if not plugin_file:
346            continue
347
348        files_to_scan = []
349        plugin_dir = os.path.dirname(plugin_file)
350        if os.path.basename(plugin_file) == '__init__.py':
351            for root, _, files in os.walk(plugin_dir):
352                for fname in files:
353                    if fname.endswith('.py'):
354                        files_to_scan.append(os.path.join(root, fname))
355        else:
356            files_to_scan.append(plugin_file)
357
358        found = False
359        for fpath in files_to_scan:
360            try:
361                with open(fpath, encoding='utf-8') as f:
362                    text = f.read()
363            except (OSError, UnicodeDecodeError):
364                continue
365            if 'make_connector' in text or 'Connector' in text:
366                found = True
367                break
368
369        if found:
370            to_import.append(plugin.name)
371
372    if not to_import:
373        return
374
375    import_plugins(*to_import)
376
377
378def unload_plugin_connectors(
379    plugin_names: Optional[List[str]] = None,
380    debug: bool = False,
381) -> None:
382    """
383    Unload custom connectors added by plugins.
384    """
385    from meerschaum.plugins import get_plugins_names
386    global custom_types, _known_custom_types, types, plugins_types, connectors, _loaded_plugin_connectors
387
388    plugin_names = plugin_names or get_plugins_names()
389
390    for plugin_name in plugin_names:
391        plugin_types = plugins_types.get(plugin_name, [])
392        for typ in plugin_types:
393            _ = types.pop(typ, None)
394            _ = connectors.pop(typ, None)
395            if typ in instance_types:
396                instance_types.remove(typ)
397
398    custom_types.clear()
399    custom_types.update(_known_custom_types)
400    _loaded_plugin_connectors = False
401
402
403def get_connector_plugin(
404    connector: Connector,
405) -> Union[str, None, mrsm.Plugin]:
406    """
407    Determine the plugin for a connector.
408    This is useful for handling virtual environments for custom instance connectors.
409
410    Parameters
411    ----------
412    connector: Connector
413        The connector which may require a virtual environment.
414
415    Returns
416    -------
417    A Plugin, 'mrsm', or None.
418    """
419    if not hasattr(connector, 'type'):
420        return None
421    plugin_name = (
422        connector.__module__.replace('plugins.', '').split('.')[0]
423        if connector.type in custom_types else (
424            connector.label
425            if connector.type == 'plugin'
426            else 'mrsm'
427        )
428    )
429    plugin = mrsm.Plugin(plugin_name)
430    return plugin if plugin.is_installed() else None
431
432
433def _load_builtin_custom_connectors():
434    """
435    Import custom connectors decorated with `@make_connector` or `@make_executor`.
436    """
437    import meerschaum.jobs.systemd
438    import meerschaum.connectors.valkey
439    _known_custom_types.add('valkey')
440    _known_custom_types.add('systemd')