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        try:
218            yield
219        except Exception:
220            pass
221        return
222
223    from meerschaum.config import _config, set_config
224    import meerschaum.config.paths as paths
225
226    old_environ = dict(os.environ)
227    old_config = copy.deepcopy(_config())
228    old_root_dir_path = paths.ROOT_DIR_PATH
229    old_plugins_dir_paths = paths.PLUGINS_DIR_PATHS
230    old_venvs_dir_path = paths.VIRTENV_RESOURCES_PATH
231    old_config_dir_path = paths.CONFIG_DIR_PATH
232
233    root_dir_env_var = STATIC_CONFIG['environment']['root']
234    plugins_dir_env_var = STATIC_CONFIG['environment']['plugins']
235    config_dir_env_var = STATIC_CONFIG['environment']['config_dir']
236    venvs_dir_env_var = STATIC_CONFIG['environment']['venvs']
237
238    replaced_root = False
239    replaced_plugins = False
240    replaced_venvs = False
241    replaced_config_dir = False
242
243    ### Hold the lock only while swapping the process-global state, not across the
244    ### `yield` (which runs arbitrary user code and could deadlock or serialize the
245    ### whole daemon).
246    with _replace_env_lock:
247        os.environ.update(env)
248
249        if root_dir_env_var in env:
250            root_dir_path = pathlib.Path(env[root_dir_env_var])
251            paths.set_root(root_dir_path)
252            replaced_root = True
253
254        if plugins_dir_env_var in env:
255            plugins_dir_paths = env[plugins_dir_env_var]
256            paths.set_plugins_dir_paths(plugins_dir_paths)
257            replaced_plugins = True
258
259        if venvs_dir_env_var in env:
260            venv_dir_path = pathlib.Path(env[venvs_dir_env_var])
261            paths.set_venvs_dir_path(venv_dir_path)
262            replaced_venvs = True
263
264        if config_dir_env_var in env:
265            config_dir_path = pathlib.Path(env[config_dir_env_var])
266            paths.set_config_dir_path(config_dir_path)
267            replaced_config_dir = True
268
269        apply_environment_patches(env)
270        apply_environment_uris(env)
271
272        ### The `plugins` package in `sys.modules` caches a `__path__` resolved
273        ### against the previous scope's `PLUGINS_RESOURCES_PATH`. If the root or
274        ### plugins-dir scope actually changed, drop that cache so the next plugin
275        ### import re-discovers plugins under the new scope (otherwise module-level
276        ### `from_plugin_import` of a sibling plugin fails with
277        ### "Unable to import plugin '<name>'" and connector-providing plugins
278        ### don't register their connectors).
279        plugins_scope_changed = (
280            (replaced_root and paths.ROOT_DIR_PATH != old_root_dir_path)
281            or (replaced_plugins and paths.PLUGINS_DIR_PATHS != old_plugins_dir_paths)
282        )
283        if plugins_scope_changed:
284            from meerschaum.plugins import invalidate_plugins_cache
285            invalidate_plugins_cache()
286
287    try:
288        yield
289    finally:
290        with _replace_env_lock:
291            ### Restore `os.environ` by diffing, NOT `clear()` + `update()`.
292            ### `clear()` blanks the entire environment for a window during which a
293            ### concurrent thread (e.g. a daemon's `sync_plugins_symlinks`) sees
294            ### `MRSM_PLUGINS_DIR` as absent and falls back to the host plugins dir,
295            ### so it never creates a project's plugin symlinks (the `plugin:<name>`
296            ### job then fails to import). Keys present in `old_environ` are never
297            ### removed, so vars like `MRSM_PLUGINS_DIR` never momentarily vanish.
298            for key in [k for k in os.environ if k not in old_environ]:
299                del os.environ[key]
300            for key, val in old_environ.items():
301                if os.environ.get(key) != val:
302                    os.environ[key] = val
303
304            if replaced_root:
305                paths.set_root(old_root_dir_path)
306
307            if replaced_plugins:
308                paths.set_plugins_dir_paths(old_plugins_dir_paths)
309
310            if replaced_venvs:
311                paths.set_venvs_dir_path(old_venvs_dir_path)
312
313            if replaced_config_dir:
314                paths.set_config_dir_path(old_config_dir_path)
315
316            _config().clear()
317            set_config(old_config)
318
319            ### Mirror the invalidation on exit: the scope is being restored,
320            ### so the `plugins` package cached during the temporary scope is
321            ### now the stale one.
322            if plugins_scope_changed:
323                from meerschaum.plugins import invalidate_plugins_cache
324                invalidate_plugins_cache()