[Issue 1135] Setup lookup value logic within the API

api/src/adapters/db/type_decorators/postgres_type_decorators.py

@@ -1,74 +1,47 @@
-from enum import StrEnum
 from typing import Any, Type
 
-from sqlalchemy import Text
+from sqlalchemy import Integer
 from sqlalchemy.types import TypeDecorator
 
+from src.db.models.lookup import LookupRegistry, LookupTable
 
-class StrEnumColumn(TypeDecorator):
-    """
-    This class handles converting StrEnum objects into strings when writing to the DB,
-    and converting those strings back into the provided StrEnum when reading from the DB.
-
-    Example Usage::
-
-        from enum import StrEnum
-        from sqlalchemy.orm import Mapped, mapped_column
-        from src.db.models.base import Base, IdMixin, TimestampMixin
-
-        # Define a StrEnum somewhere
-        class ExampleEnum(StrEnum):
-            VALUE_A = "a"
-            VALUE_B = "b"
-            VALUE_C = "c"
-
-        # Create your DB model, specifying the column type like so
-        class Example(Base, IdMixin, TimestampMixin):
-            __tablename__ = "example"
-
-            example_column: Mapped[ExampleEnum] = mapped_column(StrEnumColumn(ExampleEnum))
-            ...
-
-
-        # Use the model - when the value is written to the DB, just "a" will be written
-        example = Example(example_column=ExampleEnum.VALUE_A)
 
+class LookupColumn(TypeDecorator):
+    """
+    A Postgres column decorator that wraps
+    an integer column representing a lookup int.
 
-    See: https://docs.sqlalchemy.org/en/20/core/custom_types.html#types-typedecorator
+    This takes in the LookupTable that the lookup value
+    is stored in, and handles converting the Lookup object
+    in-code into the integer in the DB automatically (and the reverse).
     """
 
-    impl = Text
+    impl = Integer
 
     cache_ok = True
 
-    def __init__(self, lookup_enum: Type[StrEnum], *args: Any, **kwargs: Any):
+    def __init__(self, lookup_table: Type[LookupTable], *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self.lookup_enum = lookup_enum
+        self.lookup_table = lookup_table
 
-    def process_bind_param(self, value: Any | None, dialect: Any) -> str | None:
-        """
-        Method for converting a StrEnum when writing TO the DB
-        """
+    def process_bind_param(self, value: Any | None, dialect: Any) -> int | None:
         if value is None:
             return None
 
-        if not isinstance(value, self.lookup_enum):
+        if not LookupRegistry.is_valid_type_for_table(self.lookup_table, value):
             raise Exception(
-                f"Cannot convert value of type {type(value)} for binding column, expected {self.lookup_enum}"
+                f"Cannot convert value of type {type(value)} for binding column in table {self.lookup_table.get_table_name()}"
             )
 
-        return value  # technically a StrEnum which subclasses str
+        return LookupRegistry.get_lookup_int_for_enum(self.lookup_table, value)
 
     def process_result_value(self, value: Any | None, dialect: Any) -> Any | None:
-        """
-        Method for converting a string in the DB back to the StrEnum
-        """
         if value is None:
             return None
 
-        if not isinstance(value, str):
-            raise Exception(f"Cannot process value from DB of type {type(value)}")
+        if not isinstance(value, int):
+            raise Exception(
+                f"Cannot process value from DB of type {type(value)} in table {self.lookup_table.get_table_name()}"
+            )
 
-        # This calls the constructor of the enum (eg. MyEnum(value)) and gives
-        # an instance of that enum
-        return self.lookup_enum(value)
+        return LookupRegistry.get_enum_for_lookup_int(self.lookup_table, value)

api/src/constants/lookup_constants.py

@@ -1,9 +1,23 @@
 from enum import StrEnum
 
+from src.db.models.lookup import LookupConfig, LookupStr
+
 
 class OpportunityCategory(StrEnum):
+    # TODO - change these to full text once we build the next version of the API
     DISCRETIONARY = "D"
     MANDATORY = "M"
     CONTINUATION = "C"
     EARMARK = "E"
     OTHER = "O"
+
+
+OPPORTUNITY_CATEGORY_CONFIG = LookupConfig(
+    [
+        LookupStr(OpportunityCategory.DISCRETIONARY, 1),
+        LookupStr(OpportunityCategory.MANDATORY, 2),
+        LookupStr(OpportunityCategory.CONTINUATION, 3),
+        LookupStr(OpportunityCategory.EARMARK, 4),
+        LookupStr(OpportunityCategory.OTHER, 5),
+    ]
+)

api/src/db/migrations/env.py

@@ -6,9 +6,10 @@
 
 import src.adapters.db as db
 import src.logging
-from src.adapters.db.type_decorators.postgres_type_decorators import StrEnumColumn
 from src.db.models import metadata
 
+from src.adapters.db.type_decorators.postgres_type_decorators import LookupColumn  # isort:skip
+
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
 config = context.config
@@ -41,11 +42,11 @@ def include_object(
             return True
 
     def render_item(type_: str, obj: Any, autogen_context: Any) -> Any:
-        # Alembic tries to set the type of the column as StrEnumColumn
-        # despite it being derived from the Text column type,
-        # so force it to be Text during it's generation process
-        if type_ == "type" and isinstance(obj, StrEnumColumn):
-            return "sa.Text()"
+        # Alembic tries to set the type of the column as LookupColumn
+        # despite it being derived from the Integer column type,
+        # so force it to be Integer during it's generation process
+        if type_ == "type" and isinstance(obj, LookupColumn):
+            return "sa.Integer()"
 
         # False means to use the default processing
         return False

api/src/db/migrations/run.py

@@ -10,6 +10,9 @@
 from alembic.config import Config
 from alembic.runtime import migration
 
+import src.logging
+from src.db.models.lookup.sync_lookup_values import sync_lookup_values
+
 logger = logging.getLogger(__name__)
 alembic_cfg = Config(os.path.join(os.path.dirname(__file__), "./alembic.ini"))
 
@@ -20,6 +23,11 @@
 def up(revision: str = "head") -> None:
     command.upgrade(alembic_cfg, revision)
 
+    # We want logging for the lookups, but alembic already sets
+    # it up in env.py, so set it up again separately for the syncing
+    with src.logging.init("sync_lookup_values"):
+        sync_lookup_values()
+
 
 def down(revision: str = "-1") -> None:
     command.downgrade(alembic_cfg, revision)

api/src/db/migrations/versions/2024_02_02_add_opportunity_category_table.py

@@ -0,0 +1,73 @@
+"""add opportunity category table
+
+Revision ID: 479221fb8ba8
+Revises: b1eb1bd4a647
+Create Date: 2024-02-02 11:36:33.241412
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "479221fb8ba8"
+down_revision = "b1eb1bd4a647"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "lk_opportunity_category",
+        sa.Column("opportunity_category_id", sa.Integer(), nullable=False),
+        sa.Column("description", sa.Text(), nullable=False),
+        sa.Column(
+            "created_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            sa.TIMESTAMP(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.PrimaryKeyConstraint(
+            "opportunity_category_id", name=op.f("lk_opportunity_category_pkey")
+        ),
+    )
+    op.add_column("opportunity", sa.Column("opportunity_category_id", sa.Integer(), nullable=True))
+    op.drop_index("opportunity_category_idx", table_name="opportunity")
+    op.create_index(
+        op.f("opportunity_opportunity_category_id_idx"),
+        "opportunity",
+        ["opportunity_category_id"],
+        unique=False,
+    )
+    op.create_foreign_key(
+        op.f("opportunity_opportunity_category_id_lk_opportunity_category_fkey"),
+        "opportunity",
+        "lk_opportunity_category",
+        ["opportunity_category_id"],
+        ["opportunity_category_id"],
+    )
+    op.drop_column("opportunity", "category")
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "opportunity", sa.Column("category", sa.TEXT(), autoincrement=False, nullable=True)
+    )
+    op.drop_constraint(
+        op.f("opportunity_opportunity_category_id_lk_opportunity_category_fkey"),
+        "opportunity",
+        type_="foreignkey",
+    )
+    op.drop_index(op.f("opportunity_opportunity_category_id_idx"), table_name="opportunity")
+    op.create_index("opportunity_category_idx", "opportunity", ["category"], unique=False)
+    op.drop_column("opportunity", "opportunity_category_id")
+    op.drop_table("lk_opportunity_category")
+    # ### end Alembic commands ###

api/src/db/models/__init__.py

@@ -1,6 +1,6 @@
 import logging
 
-from . import base, opportunity_models
+from . import base, lookup_models, opportunity_models
 from .transfer import topportunity_models
 
 logger = logging.getLogger(__name__)
@@ -9,4 +9,4 @@
 # This is used by tests to create the test database.
 metadata = base.metadata
 
-__all__ = ["metadata", "opportunity_models", "topportunity_models"]
+__all__ = ["metadata", "opportunity_models", "lookup_models", "topportunity_models"]

api/src/db/models/lookup/__init__.py

@@ -0,0 +1,14 @@
+from .lookup import Lookup, LookupConfig, LookupInt, LookupStr
+from .lookup_registry import LookupRegistry
+from .lookup_table import LookupTable
+from .sync_lookup_values import sync_lookup_values
+
+__all__ = [
+    "Lookup",
+    "LookupInt",
+    "LookupStr",
+    "LookupConfig",
+    "LookupTable",
+    "LookupRegistry",
+    "sync_lookup_values",
+]

api/src/db/models/lookup/lookup.py

@@ -0,0 +1,120 @@
+import dataclasses
+from abc import ABC, ABCMeta, abstractmethod
+from enum import IntEnum, StrEnum
+from typing import Generic, Optional, Tuple, Type, TypeVar
+
+T = TypeVar("T", StrEnum, IntEnum)
+
+
+@dataclasses.dataclass
+class Lookup(Generic[T], ABC, metaclass=ABCMeta):
+    """
+    A class which handles mapping a specific enum
+    member to additional metadata.
+
+    At the moment, it only specifies a lookup value in
+    the DB, but we can expand this to include configuration
+    for additional member-specific fields like whether the
+    field is deprecated, or should be excluded from the API schema
+    """
+
+    lookup_enum: T
+    lookup_val: int
+
+    @abstractmethod
+    def get_description(self) -> str:
+        pass
+
+
+class LookupStr(Lookup[StrEnum]):
+    def get_description(self) -> str:
+        return self.lookup_enum
+
+
+class LookupInt(Lookup[IntEnum]):
+    def get_description(self) -> str:
+        return self.lookup_enum.name
+
+
+class LookupConfig(Generic[T]):
+    """
+    Configuration object for storing lookup mapping
+    information. Helps with the conversion of our
+    enums to lookup integers in the DB, and vice-versa.
+    """
+
+    _enums: Tuple[Type[T], ...]
+    _enum_to_lookup_map: dict[T, Lookup]
+    _int_to_lookup_map: dict[int, Lookup]
+
+    def __init__(self, lookups: list[Lookup]) -> None:
+        enum_types_seen: set[Type[T]] = set()
+        _enum_to_lookup_map: dict[T, Lookup] = {}
+        _int_to_lookup_map: dict[int, Lookup] = {}
+
+        for lookup in lookups:
+            if lookup.lookup_enum in _enum_to_lookup_map:
+                raise AttributeError(
+                    f"Duplicate lookup_enum {lookup.lookup_enum} defined, {lookup} + {_enum_to_lookup_map[lookup.lookup_enum]}"
+                )
+            _enum_to_lookup_map[lookup.lookup_enum] = lookup
+
+            if lookup.lookup_val <= 0:
+                raise AttributeError(
+                    f"Only positive lookup_val values are allowed, {lookup} not allowed"
+                )
+
+            if lookup.lookup_val in _int_to_lookup_map:
+                raise AttributeError(
+                    f"Duplicate lookup_val {lookup.lookup_val} defined, {lookup} + {_int_to_lookup_map[lookup.lookup_val]}"
+                )
+            _int_to_lookup_map[lookup.lookup_val] = lookup
+
+            enum_types_seen.add(lookup.lookup_enum.__class__)
+
+        # Verify that for each enum in the config
+        # that all of the values were mapped
+        expected_enum_members = set()
+        for enum_type_seen in enum_types_seen:
+            expected_enum_members.update([e for e in enum_type_seen])
+
+        diff = expected_enum_members.difference(_enum_to_lookup_map)
+        if len(diff) > 0:
+            raise AttributeError(
+                f"Lookup config must define a mapping for all enum values, the following were missing: {diff}"
+            )
+
+        self._enums: Tuple[Type[T], ...] = tuple(enum_types_seen)
+        self._enum_to_lookup_map: dict[T, Lookup] = _enum_to_lookup_map
+        self._int_to_lookup_map: dict[int, Lookup] = _int_to_lookup_map
+
+    def get_enums(self) -> Tuple[Type[T], ...]:
+        return self._enums
+
+    def get_lookups(self) -> list[Lookup]:
+        return [lk for lk in self._enum_to_lookup_map.values()]
+
+    def get_int_for_enum(self, e: T) -> Optional[int]:
+        """
+        Given an enum, get the lookup int for it in the DB
+        """
+        lookup = self._enum_to_lookup_map.get(e)
+        if lookup is None:
+            return None
+
+        return lookup.lookup_val
+
+    def get_lookup_for_int(self, num: int) -> Optional[Lookup]:
+        """
+        Given a lookup int, get the lookup for it
+        """
+        return self._int_to_lookup_map.get(num)
+
+    def get_enum_for_int(self, num: int) -> Optional[T]:
+        """
+        Given a lookup int, get the enum for it (via the lookup object)
+        """
+        lookup = self.get_lookup_for_int(num)
+        if lookup is None:
+            return None
+        return lookup.lookup_enum