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