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    from meerschaum.plugins import get_plugins, import_plugins
338    to_import = []
339    for plugin in get_plugins():
340        if plugin is None:
341            continue
342
343        with open(plugin.__file__, encoding='utf-8') as f:
344            text = f.read()
345
346        if 'make_connector' in text or 'Connector' in text:
347            to_import.append(plugin.name)
348
349    if not to_import:
350        return
351
352    import_plugins(*to_import)
353
354
355def unload_plugin_connectors(
356    plugin_names: Optional[List[str]] = None,
357    debug: bool = False,
358) -> None:
359    """
360    Unload custom connectors added by plugins.
361    """
362    from meerschaum.plugins import get_plugins_names
363    global custom_types, _known_custom_types, types, plugins_types, connectors, _loaded_plugin_connectors
364
365    plugin_names = plugin_names or get_plugins_names()
366
367    for plugin_name in plugin_names:
368        plugin_types = plugins_types.get(plugin_name, [])
369        for typ in plugin_types:
370            _ = types.pop(typ, None)
371            _ = connectors.pop(typ, None)
372            if typ in instance_types:
373                instance_types.remove(typ)
374
375    custom_types.clear()
376    custom_types.update(_known_custom_types)
377    _loaded_plugin_connectors = False
378
379
380def get_connector_plugin(
381    connector: Connector,
382) -> Union[str, None, mrsm.Plugin]:
383    """
384    Determine the plugin for a connector.
385    This is useful for handling virtual environments for custom instance connectors.
386
387    Parameters
388    ----------
389    connector: Connector
390        The connector which may require a virtual environment.
391
392    Returns
393    -------
394    A Plugin, 'mrsm', or None.
395    """
396    if not hasattr(connector, 'type'):
397        return None
398    plugin_name = (
399        connector.__module__.replace('plugins.', '').split('.')[0]
400        if connector.type in custom_types else (
401            connector.label
402            if connector.type == 'plugin'
403            else 'mrsm'
404        )
405    )
406    plugin = mrsm.Plugin(plugin_name)
407    return plugin if plugin.is_installed() else None
408
409
410def _load_builtin_custom_connectors():
411    """
412    Import custom connectors decorated with `@make_connector` or `@make_executor`.
413    """
414    import meerschaum.jobs.systemd
415    import meerschaum.connectors.valkey
416    _known_custom_types.add('valkey')
417    _known_custom_types.add('systemd')