Add TOML config file support (#8121) (#8215)

This adds a hidden server command-line argument
--config-file, as well as the "config-file" key in multi-
tenant config file.

Replaces #8059, Fixes #1325, Fixes #7990

Co-authored-by: Fantix King <[email protected]>

edb/common/asyncutil.py

@@ -18,9 +18,20 @@
 
 
 from __future__ import annotations
-from typing import Callable, TypeVar, Awaitable
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    cast,
+    overload,
+    Self,
+    TypeVar,
+    Type,
+)
 
 import asyncio
+import inspect
+import warnings
 
 
 _T = TypeVar('_T')
@@ -140,3 +151,188 @@ async def debounce(
         batch = []
         last_signal = t
         target_time = None
+
+
+_Owner = TypeVar("_Owner")
+HandlerFunction = Callable[[], Awaitable[None]]
+HandlerMethod = Callable[[Any], Awaitable[None]]
+
+
+class ExclusiveTask:
+    """Manages to run a repeatable task once at a time."""
+
+    _handler: HandlerFunction
+    _task: asyncio.Task | None
+    _scheduled: bool
+    _stop_requested: bool
+
+    def __init__(self, handler: HandlerFunction) -> None:
+        self._handler = handler
+        self._task = None
+        self._scheduled = False
+        self._stop_requested = False
+
+    @property
+    def scheduled(self) -> bool:
+        return self._scheduled
+
+    async def _run(self) -> None:
+        if self._scheduled and not self._stop_requested:
+            self._scheduled = False
+        else:
+            return
+        try:
+            await self._handler()
+        finally:
+            if self._scheduled and not self._stop_requested:
+                self._task = asyncio.create_task(self._run())
+            else:
+                self._task = None
+
+    def schedule(self) -> None:
+        """Schedule to run the task as soon as possible.
+
+        If already scheduled, nothing happens; it won't queue up.
+
+        If the task is already running, it will be scheduled to run again as
+        soon as the running task is done.
+        """
+        if not self._stop_requested:
+            self._scheduled = True
+            if self._task is None:
+                self._task = asyncio.create_task(self._run())
+
+    async def stop(self) -> None:
+        """Cancel scheduled task and wait for the running one to finish.
+
+        After an ExclusiveTask is stopped, no more new schedules are allowed.
+        Note: "cancel scheduled task" only means setting self._scheduled to
+        False; if an asyncio task is scheduled, stop() will still wait for it.
+        """
+        self._scheduled = False
+        self._stop_requested = True
+        if self._task is not None:
+            await self._task
+
+
+class ExclusiveTaskProperty:
+    _method: HandlerMethod
+    _name: str | None
+
+    def __init__(
+        self, method: HandlerMethod, *, slot: str | None = None
+    ) -> None:
+        self._method = method
+        self._name = slot
+
+    def __set_name__(self, owner: Type[_Owner], name: str) -> None:
+        if (slots := getattr(owner, "__slots__", None)) is not None:
+            if self._name is None:
+                raise TypeError("missing slot in @exclusive_task()")
+            if self._name not in slots:
+                raise TypeError(
+                    f"slot {self._name!r} must be defined in __slots__"
+                )
+
+        if self._name is None:
+            self._name = name
+
+    @overload
+    def __get__(self, instance: None, owner: Type[_Owner]) -> Self: ...
+
+    @overload
+    def __get__(
+        self, instance: _Owner, owner: Type[_Owner]
+    ) -> ExclusiveTask: ...
+
+    def __get__(
+        self, instance: _Owner | None, owner: Type[_Owner]
+    ) -> ExclusiveTask | Self:
+        # getattr on the class
+        if instance is None:
+            return self
+
+        assert self._name is not None
+
+        # getattr on an object with __dict__
+        if (d := getattr(instance, "__dict__", None)) is not None:
+            if rv := d.get(self._name, None):
+                return rv
+            rv = ExclusiveTask(self._method.__get__(instance, owner))
+            d[self._name] = rv
+            return rv
+
+        # getattr on an object with __slots__
+        else:
+            if rv := getattr(instance, self._name, None):
+                return rv
+            rv = ExclusiveTask(self._method.__get__(instance, owner))
+            setattr(instance, self._name, rv)
+            return rv
+
+
+ExclusiveTaskDecorator = Callable[
+    [HandlerFunction | HandlerMethod], ExclusiveTask | ExclusiveTaskProperty
+]
+
+
+def _exclusive_task(
+    handler: HandlerFunction | HandlerMethod, *, slot: str | None
+) -> ExclusiveTask | ExclusiveTaskProperty:
+    sig = inspect.signature(handler)
+    params = list(sig.parameters.values())
+    if len(params) == 0:
+        handler = cast(HandlerFunction, handler)
+        if slot is not None:
+            warnings.warn(
+                "slot is specified but unused in @exclusive_task()",
+                stacklevel=2,
+            )
+        return ExclusiveTask(handler)
+    elif len(params) == 1 and params[0].kind in (
+        inspect.Parameter.POSITIONAL_ONLY,
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+    ):
+        handler = cast(HandlerMethod, handler)
+        return ExclusiveTaskProperty(handler, slot=slot)
+    else:
+        raise TypeError("bad signature")
+
+
+@overload
+def exclusive_task(handler: HandlerFunction) -> ExclusiveTask: ...
+
+
+@overload
+def exclusive_task(
+    handler: HandlerMethod, *, slot: str | None = None
+) -> ExclusiveTaskProperty: ...
+
+
+@overload
+def exclusive_task(*, slot: str | None = None) -> ExclusiveTaskDecorator: ...
+
+
+def exclusive_task(
+    handler: HandlerFunction | HandlerMethod | None = None,
+    *,
+    slot: str | None = None,
+) -> ExclusiveTask | ExclusiveTaskProperty | ExclusiveTaskDecorator:
+    """Convert an async function into an ExclusiveTask.
+
+    This decorator can be applied to either top-level functions or methods
+    in a class. In the latter case, the exclusiveness is bound to each object
+    of the owning class. If the owning class defines __slots__, you must also
+    define an extra slot to store the exclusive state and tell exclusive_task()
+    by providing the `slot` argument.
+    """
+    if handler is None:
+
+        def decorator(
+            handler: HandlerFunction | HandlerMethod,
+        ) -> ExclusiveTask | ExclusiveTaskProperty:
+            return _exclusive_task(handler, slot=slot)
+
+        return decorator
+
+    return _exclusive_task(handler, slot=slot)

edb/pgsql/metaschema.py

@@ -3580,6 +3580,32 @@ class SysConfigFullFunction(trampoline.VersionedFunction):
             SELECT * FROM config_defaults WHERE name like '%::%'
         ),
 
+        config_static AS (
+            SELECT
+                s.name AS name,
+                s.value AS value,
+                (CASE
+                    WHEN s.type = 'A' THEN 'command line'
+                    -- Due to inplace upgrade limits, without adding a new
+                    -- layer, configuration file values are manually squashed
+                    -- into the `environment variables` layer, see below.
+                    ELSE 'environment variable'
+                END) AS source,
+                config_spec.backend_setting IS NOT NULL AS is_backend
+            FROM
+                _edgecon_state s
+                INNER JOIN config_spec ON (config_spec.name = s.name)
+            WHERE
+                -- Give precedence to configuration file values over
+                -- environment variables manually.
+                s.type = 'A' OR s.type = 'F' OR (
+                    s.type = 'E' AND NOT EXISTS (
+                        SELECT 1 FROM _edgecon_state ss
+                        WHERE ss.name = s.name AND ss.type = 'F'
+                    )
+                )
+        ),
+
         config_sys AS (
             SELECT
                 s.key AS name,
@@ -3610,16 +3636,12 @@ class SysConfigFullFunction(trampoline.VersionedFunction):
             SELECT
                 s.name AS name,
                 s.value AS value,
-                (CASE
-                    WHEN s.type = 'A' THEN 'command line'
-                    WHEN s.type = 'E' THEN 'environment variable'
-                    ELSE 'session'
-                END) AS source,
-                FALSE AS from_backend  -- only 'B' is for backend settings
+                'session' AS source,
+                FALSE AS is_backend  -- only 'B' is for backend settings
             FROM
                 _edgecon_state s
             WHERE
-                s.type != 'B'
+                s.type = 'C'
         ),
 
         pg_db_setting AS (
@@ -3789,6 +3811,7 @@ class SysConfigFullFunction(trampoline.VersionedFunction):
             FROM
                 (
                     SELECT * FROM config_defaults UNION ALL
+                    SELECT * FROM config_static UNION ALL
                     SELECT * FROM config_sys UNION ALL
                     SELECT * FROM config_db UNION ALL
                     SELECT * FROM config_sess

edb/pgsql/patches.py

@@ -77,11 +77,12 @@ def _setup_patches(patches: list[tuple[str, str]]) -> list[tuple[str, str]]:
  * sql-introspection - refresh all sql introspection views
 """
 PATCHES: list[tuple[str, str]] = _setup_patches([
-    # 6.0b2?
+    # 6.0b2
     # One of the sql-introspection's adds a param with a default to
     # uuid_to_oid, so we need to drop the original to avoid ambiguity.
     ('sql', '''
 drop function if exists edgedbsql_v6_2f20b3fed0.uuid_to_oid(uuid) cascade
 '''),
     ('sql-introspection', ''),
+    ('metaschema-sql', 'SysConfigFullFunction'),
 ])

edb/server/args.py

@@ -150,6 +150,7 @@ class ReloadTrigger(enum.StrEnum):
     3. Multi-tenant config file (server config)
     4. Readiness state (server or tenant config)
     5. JWT sub allowlist and revocation list (server or tenant config)
+    6. The TOML config file (server or tenant config)
     """
 
     Default = "default"
@@ -265,6 +266,7 @@ class ServerConfig(NamedTuple):
     disable_dynamic_system_config: bool
     reload_config_files: ReloadTrigger
     net_worker_mode: NetWorkerMode
+    config_file: Optional[pathlib.Path]
 
     startup_script: Optional[StartupScript]
     status_sinks: List[Callable[[str], None]]
@@ -1106,6 +1108,13 @@ def resolve_envvar_value(self, ctx: click.Context):
         default='default',
         help='Controls how the std::net workers work.',
     ),
+    click.option(
+        "--config-file", type=PathPath(), metavar="PATH",
+        envvar="GEL_SERVER_CONFIG_FILE",
+        cls=EnvvarResolver,
+        help='Path to a TOML file to configure the server.',
+        hidden=True,
+    ),
 ])
 
 
@@ -1534,6 +1543,7 @@ def parse_args(**kwargs: Any):
             "readiness_state_file",
             "jwt_sub_allowlist_file",
             "jwt_revocation_list_file",
+            "config_file",
         ):
             if kwargs.get(name):
                 opt = "--" + name.replace("_", "-")

edb/server/config/__init__.py

@@ -18,7 +18,9 @@
 
 
 from __future__ import annotations
-from typing import Any, Mapping
+from typing import Any, Mapping, TypedDict
+
+import enum
 
 import immutables
 
@@ -50,9 +52,25 @@
     'load_ext_settings_from_schema',
     'get_compilation_config',
     'QueryCacheMode',
+    'ConState', 'ConStateType',
 )
 
 
+# See edb/server/pgcon/connect.py for documentation of the types
+class ConStateType(enum.StrEnum):
+    session_config = "C"
+    backend_session_config = "B"
+    command_line_argument = "A"
+    environment_variable = "E"
+    config_file = "F"
+
+
+class ConState(TypedDict):
+    name: str
+    value: Any
+    type: ConStateType
+
+
 def lookup(
     name: str,
     *configs: Mapping[str, SettingValue],