fix: Create FlexibleEnumMeta enum metaclass (#74)

Author: serramatutu

.changes/unreleased/Fixes-20250327-184148.yaml

@@ -0,0 +1,3 @@
+kind: Fixes
+body: Added `FlexibleEnumMeta` to make enums non-breaking when we add a new value to the API.
+time: 2025-03-27T18:41:48.001433+01:00

dbtsl/models/base.py

@@ -2,9 +2,10 @@
 import warnings
 from dataclasses import dataclass, fields, is_dataclass
 from dataclasses import field as dc_field
+from enum import EnumMeta
 from functools import cache
 from types import MappingProxyType
-from typing import Any, ClassVar, Dict, List, Set, Type, Union
+from typing import Any, ClassVar, Dict, List, Set, Tuple, Type, Union
 from typing import get_args as get_type_args
 from typing import get_origin as get_type_origin
 
@@ -20,6 +21,35 @@ def snake_case_to_camel_case(s: str) -> str:
     return tokens[0] + "".join(t.title() for t in tokens[1:])
 
 
+class FlexibleEnumMeta(EnumMeta):
+    """Makes an Enum class not break if you provide it an unknown value."""
+
+    _subclass_registry: ClassVar[Set[str]] = set()
+
+    UNKNOWN = "UNKNOWN"
+
+    def __new__(metacls: Type["FlexibleEnumMeta"], name: str, bases: Tuple[Type], namespace: Dict[str, Any], **kwargs):
+        """Overwrite the _missing_ method of enum classes."""
+        msg = f"Class {name} needs UNKNOWN attribute with 'UNKNOWN' string value"
+        assert namespace.get("UNKNOWN", None) == "UNKNOWN", msg
+
+        metacls._subclass_registry.add(name)
+
+        newclass = super().__new__(metacls, name, bases, namespace)  # pyright: ignore[reportArgumentType]
+        setattr(newclass, "_missing_", classmethod(metacls._missing_))  # pyright: ignore[reportArgumentType]
+        return newclass
+
+    def __getitem__(cls, name: str) -> Any:
+        """Return the UNKNOWN attribute if can't find value in class."""
+        try:
+            return super().__getitem__(name)
+        except KeyError:
+            return cls.UNKNOWN
+
+    def _missing_(cls, _name: str) -> str:
+        return cls.UNKNOWN
+
+
 class BaseModel(DataClassDictMixin):
     """Base class for all serializable models.
 

dbtsl/models/dimension.py

@@ -2,13 +2,14 @@
 from enum import Enum
 from typing import List, Optional
 
-from dbtsl.models.base import BaseModel, GraphQLFragmentMixin
+from dbtsl.models.base import BaseModel, FlexibleEnumMeta, GraphQLFragmentMixin
 from dbtsl.models.time import TimeGranularity
 
 
-class DimensionType(str, Enum):
+class DimensionType(Enum, metaclass=FlexibleEnumMeta):
     """The type of a dimension."""
 
+    UNKNOWN = "UNKNOWN"
     CATEGORICAL = "CATEGORICAL"
     TIME = "TIME"
 

dbtsl/models/entity.py

@@ -2,12 +2,13 @@
 from enum import Enum
 from typing import Optional
 
-from dbtsl.models.base import BaseModel, GraphQLFragmentMixin
+from dbtsl.models.base import BaseModel, FlexibleEnumMeta, GraphQLFragmentMixin
 
 
-class EntityType(str, Enum):
+class EntityType(Enum, metaclass=FlexibleEnumMeta):
     """All supported entity types."""
 
+    UNKNOWN = "UNKNOWN"
     FOREIGN = "FOREIGN"
     NATURAL = "NATURAL"
     PRIMARY = "PRIMARY"

dbtsl/models/measure.py

@@ -2,12 +2,13 @@
 from enum import Enum
 from typing import Optional
 
-from dbtsl.models.base import BaseModel, GraphQLFragmentMixin
+from dbtsl.models.base import BaseModel, FlexibleEnumMeta, GraphQLFragmentMixin
 
 
-class AggregationType(str, Enum):
+class AggregationType(Enum, metaclass=FlexibleEnumMeta):
     """All supported aggregation functions."""
 
+    UNKNOWN = "UNKNOWN"
     SUM = "SUM"
     MIN = "MIN"
     MAX = "MAX"

dbtsl/models/metric.py

@@ -2,16 +2,17 @@
 from enum import Enum
 from typing import List, Optional
 
-from dbtsl.models.base import BaseModel, GraphQLFragmentMixin
+from dbtsl.models.base import BaseModel, FlexibleEnumMeta, GraphQLFragmentMixin
 from dbtsl.models.dimension import Dimension
 from dbtsl.models.entity import Entity
 from dbtsl.models.measure import Measure
 from dbtsl.models.time import TimeGranularity
 
 
-class MetricType(str, Enum):
+class MetricType(Enum, metaclass=FlexibleEnumMeta):
     """The type of a Metric."""
 
+    UNKNOWN = "UNKNOWN"
     SIMPLE = "SIMPLE"
     RATIO = "RATIO"
     CUMULATIVE = "CUMULATIVE"

dbtsl/models/query.py

@@ -6,14 +6,15 @@
 
 import pyarrow as pa
 
-from dbtsl.models.base import BaseModel, GraphQLFragmentMixin
+from dbtsl.models.base import BaseModel, FlexibleEnumMeta, GraphQLFragmentMixin
 
 QueryId = NewType("QueryId", str)
 
 
-class QueryStatus(str, Enum):
+class QueryStatus(Enum, metaclass=FlexibleEnumMeta):
     """All the possible states of a query."""
 
+    UNKNOWN = "UNKNOWN"
     PENDING = "PENDING"
     RUNNING = "RUNNING"
     COMPILED = "COMPILED"

dbtsl/models/saved_query.py

@@ -3,13 +3,14 @@
 from enum import Enum
 from typing import List, Optional
 
-from dbtsl.models.base import BaseModel, GraphQLFragmentMixin
+from dbtsl.models.base import BaseModel, FlexibleEnumMeta, GraphQLFragmentMixin
 from dbtsl.models.time import DatePart, TimeGranularity
 
 
-class ExportDestinationType(str, Enum):
+class ExportDestinationType(Enum, metaclass=FlexibleEnumMeta):
     """All kinds of export destinations."""
 
+    UNKNOWN = "UNKNOWN"
     TABLE = "TABLE"
     VIEW = "VIEW"
 

dbtsl/models/time.py

@@ -2,10 +2,10 @@
 
 from typing_extensions import override
 
-from dbtsl.models.base import DeprecatedMixin
+from dbtsl.models.base import DeprecatedMixin, FlexibleEnumMeta
 
 
-class TimeGranularity(str, DeprecatedMixin, Enum):
+class TimeGranularity(DeprecatedMixin, Enum, metaclass=FlexibleEnumMeta):
     """A time granularity."""
 
     @override
@@ -16,6 +16,7 @@ def _deprecation_message(cls) -> str:
             "Please just use strings to represent time grains."
         )
 
+    UNKNOWN = "UNKNOWN"
     NANOSECOND = "NANOSECOND"
     MICROSECOND = "MICROSECOND"
     MILLISECOND = "MILLISECOND"
@@ -29,9 +30,10 @@ def _deprecation_message(cls) -> str:
     YEAR = "YEAR"
 
 
-class DatePart(str, Enum):
+class DatePart(Enum, metaclass=FlexibleEnumMeta):
     """Date part."""
 
+    UNKNOWN = "UNKNOWN"
     DOY = "DOY"
     DOW = "DOW"
     DAY = "DAY"

tests/test_models.py

@@ -1,12 +1,15 @@
+import inspect
 import warnings
 from dataclasses import dataclass
 from dataclasses import field as dc_field
+from enum import Enum
 from typing import List
 
 import pytest
 from mashumaro.codecs.basic import decode
 from typing_extensions import override
 
+import dbtsl.models as ALL_EXPORTED_MODELS
 from dbtsl.api.graphql.util import normalize_query
 from dbtsl.api.shared.query_params import (
     AdhocQueryParametersStrict,
@@ -17,7 +20,7 @@
     validate_order_by,
     validate_query_parameters,
 )
-from dbtsl.models.base import BaseModel, DeprecatedMixin, GraphQLFragmentMixin
+from dbtsl.models.base import BaseModel, DeprecatedMixin, FlexibleEnumMeta, GraphQLFragmentMixin
 from dbtsl.models.base import snake_case_to_camel_case as stc
 
 
@@ -29,6 +32,47 @@ def test_snake_case_to_camel_case() -> None:
     assert stc("helloWorld") == "helloWorld"
 
 
+def test_FlexibleEnumMeta_parse_unknown_value() -> None:
+    """Make sure FlexibleEnumMeta classes parse unknown values without error."""
+
+    class EnumTest(Enum, metaclass=FlexibleEnumMeta):
+        A = "A"
+        B = "B"
+        UNKNOWN = "UNKNOWN"
+
+    assert EnumTest("A") == EnumTest.A
+    assert EnumTest("B") == EnumTest.B
+    assert EnumTest("test") == EnumTest.UNKNOWN
+
+
+def test_FlexibleEnumMeta_subclass_with_invalid_unknown_attribute() -> None:
+    """Make sure we'll raise an error whenever a flexible enum isn't declared properly."""
+    with pytest.raises(AssertionError):
+
+        class EnumTestNoUnknown(Enum, metaclass=FlexibleEnumMeta):
+            A = "A"
+
+        _ = EnumTestNoUnknown
+
+    with pytest.raises(AssertionError):
+
+        class EnumTestInvalidUnknown(Enum, metaclass=FlexibleEnumMeta):
+            A = "A"
+            UNKNOWN = "invalid_value"
+
+        _ = EnumTestInvalidUnknown
+
+
+def test_all_enum_models_are_flexible() -> None:
+    """Make sure we didn't forget to make any enum type flexible."""
+    exported_enum_classes = inspect.getmembers(
+        ALL_EXPORTED_MODELS, lambda member: (inspect.isclass(member) and issubclass(member, Enum))
+    )
+    for enum_class_name, _ in exported_enum_classes:
+        msg = f"Enum {enum_class_name} needs to have FlexibleEnumMeta metaclass."
+        assert enum_class_name in FlexibleEnumMeta._subclass_registry, msg
+
+
 def test_base_model_auto_alias() -> None:
     @dataclass
     class SubModel(BaseModel):