Fix type hints on LSP configuration classes and improve validation when loading (#2041)

asnare · web-flow · commit 7efd6dbd3ea7 · 2025-09-26T13:33:37.000Z
## Changes

This PR updates type hints on the configuration classes, and implements
proper checking of the configuration during loading. Changes include:

- We now distinguish between missing attributes and attributes that are
present but malformed.
- With the options for a dialect, the prompt is now optional instead of
stubbing with an empty string. (The prompt is not needed for the `FIXED`
option type.)
- For options, the combinations of choices/default are now dealt with
properly depending on the kind of option it is.
 - Some tests that were using invalid fixtures have been fixed.

### Caveats/things to watch out for when reviewing

During deserialisation, in many places we now handle unhappy paths via
exceptions rather than checking in advance: the latter is unambiguous
whereas its easy to have subtle issues when the check doesn't quite
correctly anticipate or correspond to the actual usage.

### Relevant implementation details

Sadly, most of this has been done by hand even though there are
libraries that do this sort of thing for us.

### Tests

- existing unit tests
diff --git a/src/databricks/labs/lakebridge/config.py b/src/databricks/labs/lakebridge/config.py
@@ -1,8 +1,9 @@
 import logging
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from enum import Enum, auto
 from pathlib import Path
-from typing import Any, Literal, cast
+from typing import Any, Literal, TypeVar, cast
 
 from databricks.labs.blueprint.installation import JsonValue
 from databricks.labs.blueprint.tui import Prompts
@@ -20,39 +21,112 @@ class LSPPromptMethod(Enum):
     CONFIRM = auto()
 
 
-@dataclass
+E = TypeVar("E", bound=Enum)
+
+
+def extract_string_field(data: Mapping[str, JsonValue], name: str) -> str:
+    """Extract a string field from the given mapping.
+
+    Parameters:
+        data: The mapping to get the string field from.
+        name: The name of the field to extract.
+    Raises:
+        ValueError: If the field is not present, not a string, or an empty string.
+    """
+    value = _maybe_extract_string_field(data, name, is_required=True)
+    if not value:
+        msg = f"Invalid '{name}' attribute, must be a non-empty string: {value}"
+        raise ValueError(msg)
+    return value
+
+
+def _maybe_extract_string_field(data: Mapping[str, JsonValue], name: str, *, is_required: bool) -> str | None:
+    # A variant of extract_string_field() with two differences:
+    #  - It allows for optional fields.
+    #  - A provided string may be empty.
+    # (This can't easily be folded into extract_string_field() because of the different return type.)
+    try:
+        value = data[name]
+    except KeyError as e:
+        if is_required:
+            raise ValueError(f"Missing '{name}' attribute in {data}") from e
+        return None
+    if not isinstance(value, str):
+        msg = f"Invalid '{name}' entry in {data}, expecting a string: {value}"
+        raise ValueError(msg)
+    return value
+
+
+def extract_enum_field(data: Mapping[str, JsonValue], name: str, enum_type: type[E]) -> E:
+    """Extract an enum field from the given mapping.
+
+    Parameters:
+        data: The mapping to get the enum field from.
+        name: The name of the field to extract.
+        enum_type: The enum type to use for parsing the value.
+    Raises:
+        ValueError: If the field is not present and no default is provided, or if it's present but not a valid enum value.
+    """
+    enum_value = extract_string_field(data, name)
+    try:
+        return enum_type[enum_value]
+    except ValueError as e:
+        valid_values = [m.name for m in enum_type]
+        msg = f"Invalid '{name}' entry in {data}, expecting one of [{', '.join(valid_values)}]: {enum_value}"
+        raise ValueError(msg) from e
+
+
+@dataclass(frozen=True)
 class LSPConfigOptionV1:
     flag: str
     method: LSPPromptMethod
-    prompt: str = ""
+    prompt: str | None = None
     choices: list[str] | None = None
-    default: Any = None
+    default: str | None = None
 
     @classmethod
-    def parse_all(cls, data: dict[str, Any]) -> dict[str, list["LSPConfigOptionV1"]]:
+    def parse_all(cls, data: dict[str, Sequence[JsonValue]]) -> dict[str, list["LSPConfigOptionV1"]]:
         return {key: list(LSPConfigOptionV1.parse(item) for item in value) for (key, value) in data.items()}
 
     @classmethod
-    def parse(cls, data: Any) -> "LSPConfigOptionV1":
+    def _extract_choices_field(cls, data: Mapping[str, JsonValue], prompt_method: LSPPromptMethod) -> list[str] | None:
+        try:
+            choices_unsafe = data["choices"]
+        except KeyError as e:
+            if prompt_method == LSPPromptMethod.CHOICE:
+                raise ValueError(f"Missing 'choices' attribute in {data}") from e
+            return None
+        if not isinstance(choices_unsafe, list) or not all(isinstance(item, str) for item in choices_unsafe):
+            msg = f"Invalid 'choices' entry in {data}, expecting a list of strings: {choices_unsafe}"
+            raise ValueError(msg)
+        return cast(list[str], choices_unsafe)
+
+    @classmethod
+    def parse(cls, data: JsonValue) -> "LSPConfigOptionV1":
         if not isinstance(data, dict):
             raise ValueError(f"Invalid transpiler config option, expecting a dict entry, got {data}")
-        flag: str = data.get("flag", "")
-        if not flag:
-            raise ValueError(f"Missing 'flag' entry in {data}")
-        method_name: str = data.get("method", "")
-        if not method_name:
-            raise ValueError(f"Missing 'method' entry in {data}")
-        method: LSPPromptMethod = cast(LSPPromptMethod, LSPPromptMethod[method_name])
-        prompt: str = data.get("prompt", "")
-        if not prompt:
-            raise ValueError(f"Missing 'prompt' entry in {data}")
-        choices = data.get("choices", [])
-        default = data.get("default", None)
-        return LSPConfigOptionV1(flag, method, prompt, choices, default)
+
+        # Field extraction is factored out mainly to ensure the complexity of this method is not too high.
+
+        flag = extract_string_field(data, "flag")
+        method = extract_enum_field(data, "method", LSPPromptMethod)
+        prompt = _maybe_extract_string_field(data, "prompt", is_required=method != LSPPromptMethod.FORCE)
+
+        optional: dict[str, Any] = {}
+        choices = cls._extract_choices_field(data, method)
+        if choices is not None:
+            optional["choices"] = choices
+
+        default = _maybe_extract_string_field(data, "default", is_required=False)
+        if default is not None:
+            optional["default"] = default
+
+        return LSPConfigOptionV1(flag, method, prompt, **optional)
 
     def prompt_for_value(self, prompts: Prompts) -> JsonValue:
         if self.method == LSPPromptMethod.FORCE:
             return self.default
+        assert self.prompt is not None
         if self.method == LSPPromptMethod.CONFIRM:
             return prompts.confirm(self.prompt)
         if self.method == LSPPromptMethod.QUESTION:
diff --git a/src/databricks/labs/lakebridge/transpiler/lsp/lsp_engine.py b/src/databricks/labs/lakebridge/transpiler/lsp/lsp_engine.py
@@ -7,10 +7,10 @@
 import shutil
 import sys
 import venv
-from collections.abc import Callable, Sequence, Mapping
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 import attrs
 import yaml
@@ -34,8 +34,9 @@
 from pygls.exceptions import FeatureRequestError
 from pygls.lsp.client import BaseLanguageClient
 
+from databricks.labs.blueprint.installation import JsonValue, RootJsonValue
 from databricks.labs.blueprint.wheels import ProductInfo
-from databricks.labs.lakebridge.config import LSPConfigOptionV1, TranspileConfig, TranspileResult
+from databricks.labs.lakebridge.config import LSPConfigOptionV1, TranspileConfig, TranspileResult, extract_string_field
 from databricks.labs.lakebridge.errors.exceptions import IllegalStateException
 from databricks.labs.lakebridge.helpers.file_utils import is_dbt_project_file, is_sql_file
 from databricks.labs.lakebridge.transpiler.transpile_engine import TranspileEngine
@@ -50,6 +51,16 @@
 logger = logging.getLogger(__name__)
 
 
+def _is_all_strings(values: Iterable[object]) -> bool:
+    """Typeguard, to check if all values in the iterable are strings."""
+    return all(isinstance(x, str) for x in values)
+
+
+def _is_all_sequences(values: Iterable[object]) -> bool:
+    """Typeguard, to check if all values in the iterable are sequences."""
+    return all(isinstance(x, Sequence) for x in values)
+
+
 @dataclass
 class _LSPRemorphConfigV1:
     name: str
@@ -58,29 +69,63 @@ class _LSPRemorphConfigV1:
     command_line: Sequence[str]
 
     @classmethod
-    def parse(cls, data: Mapping[str, Any]) -> _LSPRemorphConfigV1:
-        version = data.get("version", 0)
+    def parse(cls, data: Mapping[str, JsonValue]) -> _LSPRemorphConfigV1:
+        cls._check_version(data)
+        name = extract_string_field(data, "name")
+        dialects = cls._extract_dialects(data)
+        env_vars = cls._extract_env_vars(data)
+        command_line = cls._extract_command_line(data)
+        return _LSPRemorphConfigV1(name, dialects, env_vars, command_line)
+
+    @classmethod
+    def _check_version(cls, data: Mapping[str, JsonValue]) -> None:
+        try:
+            version = data["version"]
+        except KeyError as e:
+            raise ValueError("Missing 'version' attribute") from e
         if version != 1:
             raise ValueError(f"Unsupported transpiler config version: {version}")
-        name: str | None = data.get("name", None)
-        if not name:
-            raise ValueError("Missing 'name' entry")
-        dialects = data.get("dialects", [])
-        if len(dialects) == 0:
-            raise ValueError("Missing 'dialects' entry")
-        env_vars = data.get("environment", {})
-        command_line = data.get("command_line", [])
-        if len(command_line) == 0:
-            raise ValueError("Missing 'command_line' entry")
-        return _LSPRemorphConfigV1(name, dialects, env_vars, command_line)
+
+    @classmethod
+    def _extract_dialects(cls, data: Mapping[str, JsonValue]) -> Sequence[str]:
+        try:
+            dialects_unsafe = data["dialects"]
+        except KeyError as e:
+            raise ValueError("Missing 'dialects' attribute") from e
+        if not isinstance(dialects_unsafe, list) or not dialects_unsafe or not _is_all_strings(dialects_unsafe):
+            msg = f"Invalid 'dialects' attribute, expected a non-empty list of strings but got: {dialects_unsafe}"
+            raise ValueError(msg)
+        return cast(list[str], dialects_unsafe)
+
+    @classmethod
+    def _extract_env_vars(cls, data: Mapping[str, JsonValue]) -> Mapping[str, str]:
+        try:
+            env_vars_unsafe = data["environment"]
+            if not isinstance(env_vars_unsafe, Mapping) or not _is_all_strings(env_vars_unsafe.values()):
+                msg = f"Invalid 'environment' entry, expected a mapping with string values but got: {env_vars_unsafe}"
+                raise ValueError(msg)
+            return cast(dict[str, str], env_vars_unsafe)
+        except KeyError:
+            return {}
+
+    @classmethod
+    def _extract_command_line(cls, data: Mapping[str, JsonValue]) -> Sequence[str]:
+        try:
+            command_line = data["command_line"]
+        except KeyError as e:
+            raise ValueError("Missing 'command_line' attribute") from e
+        if not isinstance(command_line, list) or not command_line or not _is_all_strings(command_line):
+            msg = f"Invalid 'command_line' attribute, expected a non-empty list of strings but got: {command_line}"
+            raise ValueError(msg)
+        return cast(list[str], command_line)
 
 
 @dataclass
 class LSPConfig:
     path: Path
     remorph: _LSPRemorphConfigV1
     options: Mapping[str, Sequence[LSPConfigOptionV1]]
-    custom: Mapping[str, Any]
+    custom: Mapping[str, JsonValue]
 
     @property
     def name(self):
@@ -92,20 +137,52 @@ def options_for_dialect(self, source_dialect: str) -> Sequence[LSPConfigOptionV1
     @classmethod
     def load(cls, path: Path) -> LSPConfig:
         yaml_text = path.read_text()
-        data = yaml.safe_load(yaml_text)
-        if not isinstance(data, dict):
-            raise ValueError(f"Invalid transpiler config, expecting a dict, got a {type(data).__name__}")
-        remorph_data = data.get("remorph", None)
-        if not isinstance(remorph_data, dict):
-            raise ValueError(f"Invalid transpiler config, expecting a 'remorph' dict entry, got {remorph_data}")
-        remorph = _LSPRemorphConfigV1.parse(remorph_data)
-        options_data = data.get("options", {})
-        if not isinstance(options_data, dict):
-            raise ValueError(f"Invalid transpiler config, expecting an 'options' dict entry, got {options_data}")
-        options = LSPConfigOptionV1.parse_all(options_data)
-        custom = data.get("custom", {})
+        data: RootJsonValue = yaml.safe_load(yaml_text)
+        if not isinstance(data, Mapping):
+            msg = f"Invalid transpiler configuration, expecting a root object but got: {data}"
+            raise ValueError(msg)
+
+        remorph = cls._extract_remorph_data(data)
+        options = cls._extract_options(data)
+        custom = cls._extract_custom(data)
         return LSPConfig(path, remorph, options, custom)
 
+    @classmethod
+    def _extract_remorph_data(cls, data: Mapping[str, JsonValue]) -> _LSPRemorphConfigV1:
+        try:
+            remorph_data = data["remorph"]
+        except KeyError as e:
+            raise ValueError("Missing 'remorph' attribute") from e
+        if not isinstance(remorph_data, Mapping):
+            msg = f"Invalid transpiler config, 'remorph' entry must be an object but got: {remorph_data}"
+            raise ValueError(msg)
+        return _LSPRemorphConfigV1.parse(remorph_data)
+
+    @classmethod
+    def _extract_options(cls, data: Mapping[str, JsonValue]) -> Mapping[str, Sequence[LSPConfigOptionV1]]:
+        try:
+            options_data_unsfe = data["options"]
+        except KeyError:
+            # Optional, so no problem if missing
+            return {}
+        if not isinstance(options_data_unsfe, Mapping) or not _is_all_sequences(options_data_unsfe.values()):
+            msg = f"Invalid transpiler config, 'options' must be an object with list properties but got: {options_data_unsfe}"
+            raise ValueError(msg)
+        options_data = cast(dict[str, Sequence[JsonValue]], options_data_unsfe)
+        return LSPConfigOptionV1.parse_all(options_data)
+
+    @classmethod
+    def _extract_custom(cls, data: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
+        try:
+            custom = data["custom"]
+            if not isinstance(custom, Mapping):
+                msg = f"Invalid 'custom' entry, expected a mapping but got: {custom}"
+                raise ValueError(msg)
+            return custom
+        except KeyError:
+            # Optional, so no problem if missing
+            return {}
+
 
 def lsp_feature(name: str, options: Any | None = None):
     def wrapped(func: Callable):
diff --git a/tests/integration/transpile/test_repository.py b/tests/integration/transpile/test_repository.py
@@ -14,7 +14,6 @@ def test_installed_transpiler_info(transpiler_repository: TranspilerRepository)
         "overrides-file",
         LSPPromptMethod.QUESTION,
         "Specify the config file to override the default[Bladebridge] config - press <enter> for none",
-        [],
         default='<none>',
     )
     target_tech = LSPConfigOptionV1(
diff --git a/tests/unit/test_install.py b/tests/unit/test_install.py
@@ -1131,7 +1131,7 @@ def test_runs_and_stores_force_config_option(
     )
 
     transpiler_repository = _StubTranspilerRepository(
-        tmp_path / "labs", config_options=(LSPConfigOptionV1(flag="-XX", method=LSPPromptMethod.FORCE, default=1254),)
+        tmp_path / "labs", config_options=(LSPConfigOptionV1(flag="-XX", method=LSPPromptMethod.FORCE, default="1254"),)
     )
 
     workspace_installer = ws_installer(
@@ -1150,7 +1150,7 @@ def test_runs_and_stores_force_config_option(
     expected_config = LakebridgeConfiguration(
         transpile=TranspileConfig(
             transpiler_config_path=PATH_TO_TRANSPILER_CONFIG,
-            transpiler_options={"-XX": 1254},
+            transpiler_options={"-XX": "1254"},
             source_dialect="snowflake",
             input_source="/tmp/queries/snow",
             output_folder="/tmp/queries/databricks",
@@ -1165,7 +1165,7 @@ def test_runs_and_stores_force_config_option(
         "config.yml",
         {
             "transpiler_config_path": PATH_TO_TRANSPILER_CONFIG,
-            "transpiler_options": {'-XX': 1254},
+            "transpiler_options": {'-XX': "1254"},
             "catalog_name": "remorph_test",
             "input_source": "/tmp/queries/snow",
             "output_folder": "/tmp/queries/databricks",
diff --git a/tests/unit/transpiler/test_lsp_config.py b/tests/unit/transpiler/test_lsp_config.py

Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,6 @@ def test_installed_transpiler_info(transpiler_repository: TranspilerRepository)`
`14`	`14`	`"overrides-file",`
`15`	`15`	`LSPPromptMethod.QUESTION,`
`16`	`16`	`"Specify the config file to override the default[Bladebridge] config - press <enter> for none",`
`17`		`- [],`
`18`	`17`	`default='<none>',`
`19`	`18`	`)`
`20`	`19`	`target_tech = LSPConfigOptionV1(`