meerschaum.connectors.poll

  1#! /usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3# vim:fenc=utf-8
  4
  5"""
  6Poll database and API connections.
  7"""
  8
  9from meerschaum.utils.typing import InstanceConnector, Union, Optional, Dict, Any
 10
 11def retry_connect(
 12    connector: Union[InstanceConnector, None] = None,
 13    max_retries: int = 50,
 14    retry_wait: int = 3,
 15    workers: int = 1,
 16    warn: bool = True,
 17    print_on_connect: bool = False,
 18    enforce_chaining: bool = True,
 19    enforce_login: bool = True,
 20    debug: bool = False,
 21) -> bool:
 22    """
 23    Keep trying to connect to the database.
 24
 25    Parameters
 26    ----------
 27    connector: Union[InstanceConnector, None], default None
 28        The connector to the instance.
 29
 30    max_retries: int, default 40
 31        How many time to try connecting.
 32
 33    retry_wait: int, default 3
 34        The number of seconds between retries.
 35
 36    workers: int, default 1
 37        How many worker thread connections to make.
 38
 39    warn: bool, default True
 40        If `True`, print a warning in case the connection fails.
 41
 42    print_on_connect: bool, default False
 43        If `True`, print a message when a successful connection is established.
 44
 45    enforce_chaining: bool, default True
 46        If `False`, ignore the configured chaining option.
 47
 48    enforce_login: bool, default True
 49        If `False`, ignore an invalid login.
 50
 51    debug: bool, default False
 52        Verbosity toggle.
 53
 54    Returns
 55    -------
 56    Whether a connection could be made.
 57
 58    """
 59    import json
 60    from meerschaum.utils.venv import venv_exec
 61    from meerschaum.utils.packages import attempt_import
 62
 63    kw = {
 64        'connector_meta': connector.meta,
 65        'max_retries': max_retries,
 66        'retry_wait': retry_wait,
 67        'workers': workers,
 68        'warn': warn,
 69        'print_on_connect': print_on_connect,
 70        'enforce_chaining': enforce_chaining,
 71        'enforce_login': enforce_login,
 72        'debug': debug,
 73    }
 74
 75    dill = attempt_import('dill', lazy=False)
 76    code = (
 77        "import sys, json\n"
 78        + "from meerschaum.utils.typing import Optional, Dict, Any\n\n"
 79        + dill.source.getsource(_wrap_retry_connect) + '\n\n'
 80        + f"kw = json.loads({json.dumps(json.dumps(kw))})\n"
 81        + "success = _wrap_retry_connect(**kw)\n"
 82        + "sys.exit((0 if success else 1))"
 83    )
 84    return venv_exec(code, venv=None, debug=debug)
 85
 86
 87def _wrap_retry_connect(
 88    connector_meta: Dict[str, Any],
 89    max_retries: int = 50,
 90    retry_wait: int = 3,
 91    workers: int = 1,
 92    print_on_connect: bool = False,
 93    warn: bool = True,
 94    enforce_chaining: bool = True,
 95    enforce_login: bool = True,
 96    debug: bool = False,
 97) -> bool:
 98    """
 99    Keep trying to connect to the database.
100
101    Parameters
102    ----------
103    connector_keys: Optional[str], default None
104        The keys of the connector to the instance.
105
106    max_retries: int, default 40
107        How many time to try connecting.
108
109    retry_wait: int, default 3
110        The number of seconds between retries.
111
112    workers: int, default 1
113        How many worker thread connections to make.
114
115    warn: bool, default True
116        If `True`, print a warning in case the connection fails.
117
118    print_on_connect: bool, default False
119        If `True`, print a message when a successful connection is established.
120
121    enforce_chaining: bool, default True
122        If `False`, ignore the configured chaining option.
123
124    enforce_login: bool, default True
125        If `False`, ignore an invalid login.
126
127    debug: bool, default False
128        Verbosity toggle.
129
130    Returns
131    -------
132    Whether a connection could be made.
133
134    """
135    from meerschaum.utils.warnings import warn as _warn, error, info
136    from meerschaum.utils.debug import dprint
137    from meerschaum.connectors import instance_types, get_connector
138    from meerschaum.connectors.parse import parse_instance_keys
139    from meerschaum.utils.packages import attempt_import
140    from meerschaum.utils.misc import timed_input
141    from functools import partial
142    import platform
143    import time
144
145    connector = get_connector(**connector_meta)
146
147    if not hasattr(connector, 'test_connection'):
148        return True
149
150    retries = 0
151    connected, chaining_status = False, None
152    while retries < max_retries:
153        if debug:
154            dprint(f"Trying to connect to '{connector}'...")
155            dprint(f"Attempt ({retries + 1} / {max_retries})")
156
157        if connector.type not in ('sql', 'api'):
158            try:
159                connected = connector.test_connection()
160            except Exception as e:
161                if warn:
162                    print(e)
163                connected = False
164
165        elif connector.type == 'sql':
166
167            def _connect(_connector):
168                ### `SELECT 1` through the Rust connector (no SQLAlchemy engine).
169                if not _connector.test_connection(max_retries=1, warn=warn):
170                    raise Exception("Failed to connect.")
171
172            try:
173                _connect(connector)
174                connected = True
175            except Exception as e:
176                if warn:
177                    print(e)
178                connected = False
179
180        elif connector.type == 'api':
181            ### If the remote instance does not allow chaining, don't even try logging in.
182            if not isinstance(chaining_status, bool):
183                chaining_status = connector.get_chaining_status(debug=debug)
184                if chaining_status is None:
185                    connected = None
186                elif chaining_status is False:
187                    if enforce_chaining:
188                        if warn:
189                            _warn(
190                                f"Meerschaum instance '{connector}' does not allow chaining " +
191                                "and cannot be used as the parent for this instance.",
192                                stack=False,
193                            )
194                        return False
195
196                    ### Allow is the option to ignore chaining status.
197                    chaining_status = True
198
199            if chaining_status:
200                connected = (
201                    connector.login(warn=warn, debug=debug)[0]
202                ) if enforce_login else True
203
204                if not connected and warn:
205                    _warn(f"Unable to login to '{connector}'!", stack=False)
206
207        if connected:
208            if print_on_connect:
209                info(f"Connected to '{connector}'.")
210            return True
211
212        if warn:
213            _warn(
214                f"Connection to '{connector}' failed.\n    "
215                + f"Press [Enter] to retry or wait {retry_wait} seconds.",
216                stack = False
217            )
218            if workers and workers > 1:
219                info(
220                    f"To quit, press CTRL-C, then 'q' + [Enter] for each worker"
221                    + f" ({workers})."
222                )
223            info(f"Failed connection attempt ({retries + 1} / {max_retries})")
224
225        try:
226            if retry_wait > 0:
227                if platform.system() != 'Windows':
228                    text = timed_input(retry_wait)
229                    if text in ('q', 'quit', 'pass', 'exit', 'stop'):
230                        return None
231                else:
232                    time.sleep(retry_wait)
233        except KeyboardInterrupt:
234            return None
235        retries += 1
236
237    return False