meerschaum.config.environment

  1#! /usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3# vim:fenc=utf-8
  4
  5"""
  6Patch the runtime configuration from environment variables.
  7"""
  8
  9import os
 10import re
 11import json
 12import contextlib
 13import copy
 14import pathlib
 15import threading
 16
 17from meerschaum.utils.typing import List, Union, Dict, Any, Optional
 18from meerschaum._internal.static import STATIC_CONFIG
 19from meerschaum import _core
 20
 21### Serialize the process-global swap of `os.environ` and `meerschaum.config.paths`
 22### globals performed by `replace_env`. A daemon process runs several threads
 23### concurrently (the job target plus the `check-jobs` `RepeatTimer`), so two
 24### overlapping `replace_env` calls could otherwise leave these globals in a torn
 25### state.
 26_replace_env_lock = threading.RLock()
 27
 28
 29def apply_environment_patches(env: Optional[Dict[str, Any]] = None) -> None:
 30    """
 31    Apply patches defined in `MRSM_CONFIG` and `MRSM_PATCH`.
 32    """
 33    config_var = STATIC_CONFIG['environment']['config']
 34    patch_var = STATIC_CONFIG['environment']['patch']
 35    ### Nothing to apply: do not touch the native registry (importing `meerschaum.config` runs
 36    ### this, and a bare import must not load the extension).
 37    _env = env if env is not None else os.environ
 38    if config_var not in _env and patch_var not in _env:
 39        return
 40    return _core.apply_environment_patches(env)
 41
 42
 43def apply_environment_config(env_var: str, env: Optional[Dict[str, Any]] = None) -> None:
 44    """
 45    Parse a dictionary (simple or JSON) from an environment variable
 46    and apply it to the current configuration.
 47    """
 48    from meerschaum.config import get_config, set_config, _config
 49    from meerschaum.config._patch import apply_patch_to_config
 50
 51    env = env if env is not None else os.environ
 52
 53    if env_var not in env:
 54        return
 55
 56    from meerschaum.utils.misc import string_to_dict
 57    try:
 58        _patch = string_to_dict(str(os.environ[env_var]).lstrip())
 59    except Exception:
 60        _patch = None
 61
 62    error_msg = (
 63        f"Environment variable {env_var} is set but cannot be parsed.\n"
 64        f"Unset {env_var} or change to JSON or simplified dictionary format "
 65        "(see --help, under params for formatting)\n"
 66        f"{env_var} is set to:\n{os.environ[env_var]}\n"
 67        f"Skipping patching os environment into config..."
 68    )
 69
 70    if not isinstance(_patch, dict):
 71        print(error_msg)
 72        return
 73
 74    valids = []
 75
 76    def load_key(key: str) -> Union[Dict[str, Any], None]:
 77        try:
 78            c = get_config(key, warn=False)
 79        except Exception:
 80            c = None
 81        return c
 82
 83    ### This was multi-threaded, but I ran into all sorts of locking issues.
 84    keys = list(_patch.keys())
 85    for key in keys:
 86        _ = load_key(key)
 87
 88    ### Load and patch config files.
 89    set_config(
 90        apply_patch_to_config(
 91            _config(),
 92            _patch,
 93        )
 94    )
 95
 96
 97def apply_environment_uris(env: Optional[Dict[str, Any]] = None) -> None:
 98    """
 99    Patch temporary connectors defined in environment variables which start with
100    `MRSM_SQL_` or `MRSM_API_`.
101    """
102    env_vars = get_connector_env_vars(env=env)
103    if not env_vars:
104        return
105    return _core.apply_environment_uris(env)
106
107
108def get_connector_env_regex() -> str:
109    """
110    Return the regex pattern for valid environment variable names for instance connectors.
111    """
112    return STATIC_CONFIG['environment']['uri_regex']
113
114
115def get_connector_env_vars(env: Optional[Dict[str, Any]] = None) -> List[str]:
116    """
117    Get the names of the environment variables which match the Meerschaum connector regex.
118
119    Examples
120    --------
121    >>> get_connector_environment_vars()
122    ['MRSM_SQL_FOO']
123    """
124    uri_regex = get_connector_env_regex()
125    env_vars = []
126
127    env = env if env is not None else os.environ
128
129    for env_var in env:
130        matched = re.match(uri_regex, env_var)
131        if matched is None:
132            continue
133        if env_var in STATIC_CONFIG['environment'].values():
134            continue
135        env_vars.append(env_var)
136
137    return env_vars
138
139
140def apply_connector_uri(env_var: str, env: Optional[Dict[str, Any]] = None) -> None:
141    """
142    Parse and validate a URI obtained from an environment variable.
143    """
144    from meerschaum.config import get_config, set_config, _config
145    from meerschaum.config._patch import apply_patch_to_config
146    from meerschaum.config._read_config import search_and_substitute_config
147    from meerschaum.utils.warnings import warn
148
149    env = env if env is not None else os.environ
150
151    if env_var not in env:
152        return
153
154    uri_regex = get_connector_env_regex()
155    matched = re.match(uri_regex, env_var)
156    groups = matched.groups()
157    typ, label = groups[0].lower(), groups[1].lower()
158    if not typ or not label:
159        return
160
161    uri = env[env_var]
162
163    if uri.lstrip().startswith('{') and uri.rstrip().endswith('}'):
164        try:
165            conn_attrs = json.loads(uri)
166        except Exception:
167            warn(f"Unable to parse JSON for environment connector '{typ}:{label}'.")
168            conn_attrs = {'uri': uri}
169    else:
170        conn_attrs = {'uri': uri}
171
172    set_config(
173        apply_patch_to_config(
174            {'meerschaum': get_config('meerschaum')},
175            {'meerschaum': {'connectors': {typ: {label: conn_attrs}}}},
176        )
177    )
178
179
180def get_env_vars(env: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
181    """
182    Return all environment variables which begin with `'MRSM_'`.
183    """
184    prefix = STATIC_CONFIG['environment']['prefix']
185    env = env if env is not None else os.environ
186    return {
187        env_var: env_val
188        for env_var, env_val in env.items()
189        if env_var.startswith(prefix)
190    }
191
192
193def get_daemon_env_vars(env: Optional[Dict[str, Any]] = None) -> Dict[str, str]:
194    """
195    Return the daemon-specific environment vars in the current environment.
196    """
197    env = env if env is not None else os.environ
198
199    daemon_env_var_names = (
200        STATIC_CONFIG['environment']['systemd_log_path'],
201        STATIC_CONFIG['environment']['systemd_result_path'],
202        STATIC_CONFIG['environment']['systemd_delete_job'],
203        STATIC_CONFIG['environment']['systemd_stdin_path'],
204        STATIC_CONFIG['environment']['daemon_id'],
205    )
206    return {
207        env_var: env.get(env_var, '')
208        for env_var in daemon_env_var_names
209        if env_var in env
210    }
211
212
213@contextlib.contextmanager
214def replace_env(env: Union[Dict[str, Any], None]):
215    """
216    Temporarily replace environment variables and current configuration.
217
218    Parameters
219    ----------
220    env: Dict[str, Any]
221        The new environment dictionary to be patched on `os.environ`.
222    """
223    if env is None:
224        yield
225        return
226
227    from meerschaum.config import _config, set_config
228    import meerschaum.config.paths as paths
229    import meerschaum.utils.venv as venv_utils
230    import sys
231
232    old_environ = dict(os.environ)
233    old_config = copy.deepcopy(_config())
234    old_root_dir_path = paths.ROOT_DIR_PATH
235    old_plugins_dir_paths = paths.PLUGINS_DIR_PATHS
236    old_venvs_dir_path = paths.VIRTENV_RESOURCES_PATH
237    old_config_dir_path = paths.CONFIG_DIR_PATH
238    with venv_utils.LOCKS['active_venvs'], venv_utils.LOCKS['sys.path']:
239        old_sys_path = list(sys.path)
240        old_active_venvs = set(venv_utils.active_venvs)
241        old_active_venvs_counts = dict(venv_utils.active_venvs_counts)
242        old_active_venvs_order = list(venv_utils.active_venvs_order)
243        old_threads_active_venvs = {
244            thread_id: dict(active_counts)
245            for thread_id, active_counts in venv_utils.threads_active_venvs.items()
246        }
247
248    root_dir_env_var = STATIC_CONFIG['environment']['root']
249    plugins_dir_env_var = STATIC_CONFIG['environment']['plugins']
250    config_dir_env_var = STATIC_CONFIG['environment']['config_dir']
251    venvs_dir_env_var = STATIC_CONFIG['environment']['venvs']
252
253    replaced_root = False
254    replaced_plugins = False
255    replaced_venvs = False
256    replaced_config_dir = False
257
258    ### Hold the lock only while swapping the process-global state, not across the
259    ### `yield` (which runs arbitrary user code and could deadlock or serialize the
260    ### whole daemon).
261    with _replace_env_lock:
262        os.environ.update(env)
263
264        if root_dir_env_var in env:
265            root_dir_path = pathlib.Path(env[root_dir_env_var])
266            paths.set_root(root_dir_path)
267            replaced_root = True
268
269        if plugins_dir_env_var in env:
270            plugins_dir_paths = env[plugins_dir_env_var]
271            paths.set_plugins_dir_paths(plugins_dir_paths)
272            replaced_plugins = True
273
274        if venvs_dir_env_var in env:
275            venv_dir_path = pathlib.Path(env[venvs_dir_env_var])
276            paths.set_venvs_dir_path(venv_dir_path)
277            replaced_venvs = True
278
279        if config_dir_env_var in env:
280            config_dir_path = pathlib.Path(env[config_dir_env_var])
281            paths.set_config_dir_path(config_dir_path)
282            replaced_config_dir = True
283
284        apply_environment_patches(env)
285        apply_environment_uris(env)
286
287        ### The `plugins` package in `sys.modules` caches a `__path__` resolved
288        ### against the previous scope's `PLUGINS_RESOURCES_PATH`. If the root or
289        ### plugins-dir scope actually changed, drop that cache so the next plugin
290        ### import re-discovers plugins under the new scope (otherwise module-level
291        ### `from_plugin_import` of a sibling plugin fails with
292        ### "Unable to import plugin '<name>'" and connector-providing plugins
293        ### don't register their connectors).
294        plugins_scope_changed = (
295            (replaced_root and paths.ROOT_DIR_PATH != old_root_dir_path)
296            or (replaced_plugins and paths.PLUGINS_DIR_PATHS != old_plugins_dir_paths)
297        )
298        if plugins_scope_changed:
299            from meerschaum.plugins import invalidate_plugins_cache
300            invalidate_plugins_cache()
301
302    try:
303        yield
304    finally:
305        with _replace_env_lock:
306            ### Restore `os.environ` by diffing, NOT `clear()` + `update()`.
307            ### `clear()` blanks the entire environment for a window during which a
308            ### concurrent thread (e.g. a daemon's `sync_plugins_symlinks`) sees
309            ### `MRSM_PLUGINS_DIR` as absent and falls back to the host plugins dir,
310            ### so it never creates a project's plugin symlinks (the `plugin:<name>`
311            ### job then fails to import). Keys present in `old_environ` are never
312            ### removed, so vars like `MRSM_PLUGINS_DIR` never momentarily vanish.
313            for key in [k for k in os.environ if k not in old_environ]:
314                del os.environ[key]
315            for key, val in old_environ.items():
316                if os.environ.get(key) != val:
317                    os.environ[key] = val
318
319            if replaced_root:
320                paths.set_root(old_root_dir_path)
321
322            if replaced_plugins:
323                paths.set_plugins_dir_paths(old_plugins_dir_paths)
324
325            if replaced_venvs:
326                paths.set_venvs_dir_path(old_venvs_dir_path)
327
328            if replaced_config_dir:
329                paths.set_config_dir_path(old_config_dir_path)
330
331            ### Venv names are resolved relative to the current root. Restore both
332            ### the import path and activation bookkeeping after a scoped root swap.
333            with venv_utils.LOCKS['active_venvs'], venv_utils.LOCKS['sys.path']:
334                sys.path[:] = old_sys_path
335                venv_utils.active_venvs.clear()
336                venv_utils.active_venvs.update(old_active_venvs)
337                venv_utils.active_venvs_counts.clear()
338                venv_utils.active_venvs_counts.update(old_active_venvs_counts)
339                venv_utils.active_venvs_order[:] = old_active_venvs_order
340                venv_utils.threads_active_venvs.clear()
341                venv_utils.threads_active_venvs.update(old_threads_active_venvs)
342
343            _config().clear()
344            set_config(old_config)
345
346            ### Mirror the invalidation on exit: the scope is being restored,
347            ### so the `plugins` package cached during the temporary scope is
348            ### now the stale one.
349            if plugins_scope_changed:
350                from meerschaum.plugins import invalidate_plugins_cache
351                invalidate_plugins_cache()