🗃️(api) migrate database enum names to values (1)

jmaupetit · jmaupetit · commit 2a8ec2f70a09 · 2025-01-23T11:58:51.000+01:00
By default, python-derived enum fields are stored using the enum name
instead of its value in database. This forces us to map the
corresponding value when exporting data.

In this first step of the migration, we create the new enum types (using
values) and then alter columns to the legacy enums to a varchar.
diff --git a/src/api/CHANGELOG.md b/src/api/CHANGELOG.md
@@ -27,6 +27,7 @@ and this project adheres to
 - Cache the "get PointDeCharge id from its `id_pdc_itinerance`" database query
 - Improve JSON string parsing using pyarrow engine
 - Add default values for optional Statique model fields
+- Migrate database enum types from names to values
 - Upgrade alembic to `1.14.1`
 - Upgrade geoalchemy2 to `0.17.0`
 - Upgrade psycopg to `3.2.4`
diff --git a/src/api/qualicharge/migrations/versions/3e8ea07e5b66_alter_db_enum_to_varchar.py b/src/api/qualicharge/migrations/versions/3e8ea07e5b66_alter_db_enum_to_varchar.py
@@ -0,0 +1,232 @@
+"""Alter db enum to varchar
+
+Revision ID: 3e8ea07e5b66
+Revises: 8144a7d2553d
+Create Date: 2025-01-20 13:13:04.259534
+
+"""
+
+from enum import Enum, StrEnum
+from typing import Any, Dict, List, Sequence, Type, Union
+
+from sqlalchemy import Connection
+
+from alembic import op
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.types import VARCHAR
+from qualicharge.models.dynamic import EtatPDCEnum, EtatPriseEnum, OccupationPDCEnum
+from sqlmodel import SQLModel
+
+from qualicharge.schemas.core import PointDeCharge, Station, Status
+
+# revision identifiers, used by Alembic.
+revision: str = "3e8ea07e5b66"
+down_revision: Union[str, None] = "8144a7d2553d"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+class ImplantationStationEnum(StrEnum):
+    """Statique.implantation_station field enum."""
+
+    VOIRIE = "Voirie"
+    PARKING_PUBLIC = "Parking public"
+    PARKING_PRIVE_USAGE_PUBLIC = "Parking privé à usage public"
+    PARKING_PRIVE_CLIENTELE = "Parking privé réservé à la clientèle"
+    STATION_RECHARGE_RAPIDE = "Station dédiée à la recharge rapide"
+
+
+class ConditionAccesEnum(StrEnum):
+    """Statique.condition_acces field enum."""
+
+    ACCESS_LIBRE = "Accès libre"
+    ACCESS_RESERVE = "Accès réservé"
+
+
+class AccessibilitePMREnum(StrEnum):
+    """Statique.accessibilite_pmr field enum."""
+
+    RESERVE_PMR = "Réservé PMR"
+    NON_RESERVE = "Accessible mais non réservé PMR"
+    NON_ACCESSIBLE = "Non accessible"
+    INCONNUE = "Accessibilité inconnue"
+
+
+class RaccordementEnum(StrEnum):
+    """Statique.raccordement field enum."""
+
+    DIRECT = "Direct"
+    INDIRECT = "Indirect"
+
+
+def enum_to_dict(enum_: Type[Enum], revert: bool = False) -> Dict[str, Any]:
+    """Convert enum to dict."""
+    if revert:
+        return {member.value: member.name for member in enum_}
+    return {member.name: member.value for member in enum_}
+
+
+def migrate_db_enum(
+    connection: Connection,
+    enum_: Type[Enum],
+    existing_enum_db_name: str,
+    new_enum_db_name: str,
+    schema: SQLModel,
+    column_names: List[str],
+    revert: bool = False,
+):
+    """Migrate database Enum from keys to values (if revert if False).
+
+    When revert is True, it migrates values to keys.
+    """
+    print(f"Will migrate {enum_.__name__} for table {schema.__tablename__} ({revert=})")
+
+    # Create the new ENUM database type
+    postgresql.ENUM(
+        *enum_to_dict(enum_, revert).values(),
+        name=new_enum_db_name,
+    ).create(connection, checkfirst=True)
+
+    for column_name in column_names:
+        print(f"{column_name=}")
+
+        # Alter table column to a generic VARCHAR
+        op.alter_column(
+            schema.__tablename__,
+            column_name,
+            existing_type=postgresql.ENUM(
+                *enum_to_dict(enum_, revert).keys(),
+                name=existing_enum_db_name,
+            ),
+            type_=VARCHAR,
+            existing_nullable=False,
+            postgresql_using=f"{column_name}::VARCHAR",
+        )
+
+
+def upgrade() -> None:
+    # Alembic connection to the database
+    connection = op.get_bind()
+
+    fields_params = [
+        {
+            "enum_": ImplantationStationEnum,
+            "existing_enum_db_name": "implantationstationenum",
+            "new_enum_db_name": "implantation_station_enum",
+            "schema": Station,
+            "column_names": ["implantation_station"],
+        },
+        {
+            "enum_": ConditionAccesEnum,
+            "existing_enum_db_name": "conditionaccesenum",
+            "new_enum_db_name": "condition_acces_enum",
+            "schema": Station,
+            "column_names": ["condition_acces"],
+        },
+        {
+            "enum_": RaccordementEnum,
+            "existing_enum_db_name": "raccordementenum",
+            "new_enum_db_name": "raccordement_enum",
+            "schema": Station,
+            "column_names": ["raccordement"],
+        },
+        {
+            "enum_": AccessibilitePMREnum,
+            "existing_enum_db_name": "accessibilitepmrenum",
+            "new_enum_db_name": "accessibilite_pmr_enum",
+            "schema": PointDeCharge,
+            "column_names": ["accessibilite_pmr"],
+        },
+        {
+            "enum_": EtatPDCEnum,
+            "existing_enum_db_name": "etatpdcenum",
+            "new_enum_db_name": "etat_pdc_enum",
+            "schema": Status,
+            "column_names": ["etat_pdc"],
+        },
+        {
+            "enum_": OccupationPDCEnum,
+            "existing_enum_db_name": "occupationpdcenum",
+            "new_enum_db_name": "occupation_pdc_enum",
+            "schema": Status,
+            "column_names": ["occupation_pdc"],
+        },
+        {
+            "enum_": EtatPriseEnum,
+            "existing_enum_db_name": "etatpriseenum",
+            "new_enum_db_name": "etat_prise_enum",
+            "schema": Status,
+            "column_names": [
+                "etat_prise_type_2",
+                "etat_prise_type_combo_ccs",
+                "etat_prise_type_chademo",
+                "etat_prise_type_ef",
+            ],
+        },
+    ]
+    for field_params in fields_params:
+        migrate_db_enum(connection, **field_params)
+
+
+def downgrade() -> None:
+    # Alembic connection to the database
+    connection = op.get_bind()
+
+    fields_params = [
+        {
+            "enum_": ImplantationStationEnum,
+            "existing_enum_db_name": "implantation_station_enum",
+            "new_enum_db_name": "implantationstationenum",
+            "schema": Station,
+            "column_names": ["implantation_station"],
+        },
+        {
+            "enum_": ConditionAccesEnum,
+            "existing_enum_db_name": "condition_acces_enum",
+            "new_enum_db_name": "conditionaccesenum",
+            "schema": Station,
+            "column_names": ["condition_acces"],
+        },
+        {
+            "enum_": RaccordementEnum,
+            "existing_enum_db_name": "raccordement_enum",
+            "new_enum_db_name": "raccordementenum",
+            "schema": Station,
+            "column_names": ["raccordement"],
+        },
+        {
+            "enum_": AccessibilitePMREnum,
+            "existing_enum_db_name": "accessibilite_pmr_enum",
+            "new_enum_db_name": "accessibilitepmrenum",
+            "schema": PointDeCharge,
+            "column_names": ["accessibilite_pmr"],
+        },
+        {
+            "enum_": EtatPDCEnum,
+            "existing_enum_db_name": "etat_pdc_enum",
+            "new_enum_db_name": "etatpdcenum",
+            "schema": Status,
+            "column_names": ["etat_pdc"],
+        },
+        {
+            "enum_": OccupationPDCEnum,
+            "existing_enum_db_name": "occupation_pdc_enum",
+            "new_enum_db_name": "occupationpdcenum",
+            "schema": Status,
+            "column_names": ["occupation_pdc"],
+        },
+        {
+            "enum_": EtatPriseEnum,
+            "existing_enum_db_name": "etat_prise_enum",
+            "new_enum_db_name": "etatpriseenum",
+            "schema": Status,
+            "column_names": [
+                "etat_prise_type_2",
+                "etat_prise_type_combo_ccs",
+                "etat_prise_type_chademo",
+                "etat_prise_type_ef",
+            ],
+        },
+    ]
+    for field_params in fields_params:
+        migrate_db_enum(connection, **field_params, revert=True)
diff --git a/src/api/qualicharge/schemas/core.py b/src/api/qualicharge/schemas/core.py
@@ -20,6 +20,7 @@
 from pydantic_extra_types.coordinate import Coordinate
 from shapely.geometry import mapping
 from sqlalchemy import event
+from sqlalchemy.dialects.postgresql import ENUM as PgEnum
 from sqlalchemy.schema import Column as SAColumn
 from sqlalchemy.types import Date, DateTime, String
 from sqlmodel import Field, Relationship, UniqueConstraint, select
@@ -28,7 +29,13 @@
 
 from qualicharge.exceptions import ObjectDoesNotExist
 
-from ..models.dynamic import SessionBase, StatusBase
+from ..models.dynamic import (
+    EtatPDCEnum,
+    EtatPriseEnum,
+    OccupationPDCEnum,
+    SessionBase,
+    StatusBase,
+)
 from ..models.static import (
     AccessibilitePMREnum,
     ConditionAccesEnum,
@@ -51,6 +58,56 @@ class OperationalUnitTypeEnum(IntEnum):
     MOBILITY = 2
 
 
+# Enum definition for database: we want to store Enum values instead of keys (this is
+# the default behavior).
+def get_enum_values(enum_):
+    """Get enum values."""
+    return [m.value for m in enum_]
+
+
+ImplantationStationDBEnum: PgEnum = PgEnum(
+    ImplantationStationEnum,
+    name="implantation_station_enum",
+    values_callable=get_enum_values,
+)
+
+ConditionAccesDBEnum: PgEnum = PgEnum(
+    ConditionAccesEnum,
+    name="condition_acces_enum",
+    values_callable=get_enum_values,
+)
+
+AccessibilitePMRDBEnum: PgEnum = PgEnum(
+    AccessibilitePMREnum,
+    name="accessibilite_pmr_enum",
+    values_callable=get_enum_values,
+)
+
+RaccordementDBEnum: PgEnum = PgEnum(
+    RaccordementEnum,
+    name="raccordement_enum",
+    values_callable=get_enum_values,
+)
+
+EtatPDCDBEnum: PgEnum = PgEnum(
+    EtatPDCEnum,
+    name="etat_pdc_enum",
+    values_callable=get_enum_values,
+)
+
+EtatPriseDBEnum: PgEnum = PgEnum(
+    EtatPriseEnum,
+    name="etat_prise_enum",
+    values_callable=get_enum_values,
+)
+
+OccupationPDCDBEnum: PgEnum = PgEnum(
+    OccupationPDCEnum,
+    name="occupation_pdc_enum",
+    values_callable=get_enum_values,
+)
+
+
 class Amenageur(BaseTimestampedSQLModel, table=True):
     """Amenageur table."""
 
@@ -232,12 +289,18 @@ class Station(BaseTimestampedSQLModel, table=True):
     )
     id_station_local: Optional[str]
     nom_station: str
-    implantation_station: ImplantationStationEnum
+    implantation_station: ImplantationStationEnum = Field(
+        sa_column=SAColumn(ImplantationStationDBEnum, nullable=False)
+    )
     nbre_pdc: PositiveInt
-    condition_acces: ConditionAccesEnum
+    condition_acces: ConditionAccesEnum = Field(
+        sa_column=SAColumn(ConditionAccesDBEnum, nullable=False)
+    )
     horaires: str = Field(regex=r"(.*?)((\d{1,2}:\d{2})-(\d{1,2}:\d{2})|24/7)")
     station_deux_roues: bool
-    raccordement: Optional[RaccordementEnum]
+    raccordement: Optional[RaccordementEnum] = Field(
+        sa_column=SAColumn(RaccordementDBEnum, nullable=True)
+    )
     num_pdl: Optional[str] = Field(max_length=64)
     date_maj: NotFutureDate = Field(sa_type=Date)
     date_mise_en_service: Optional[PastDate] = Field(sa_type=Date)
@@ -325,7 +388,9 @@ class PointDeCharge(BaseTimestampedSQLModel, table=True):
     paiement_autre: Optional[bool]
     tarification: Optional[str]
     reservation: bool
-    accessibilite_pmr: AccessibilitePMREnum
+    accessibilite_pmr: AccessibilitePMREnum = Field(
+        sa_column=SAColumn(AccessibilitePMRDBEnum, nullable=False)
+    )
     restriction_gabarit: str
     observations: Optional[str]
     cable_t2_attache: Optional[bool]
@@ -379,6 +444,23 @@ class Status(BaseTimestampedSQLModel, StatusBase, table=True):
         description="The timestamp indicating when the status changed.",
     )  # type: ignore
 
+    etat_pdc: EtatPDCEnum = Field(sa_column=SAColumn(EtatPDCDBEnum, nullable=False))
+    occupation_pdc: OccupationPDCEnum = Field(
+        sa_column=SAColumn(OccupationPDCDBEnum, nullable=False)
+    )
+    etat_prise_type_2: Optional[EtatPriseEnum] = Field(
+        sa_column=SAColumn(EtatPriseDBEnum, nullable=True)
+    )
+    etat_prise_type_combo_ccs: Optional[EtatPriseEnum] = Field(
+        sa_column=SAColumn(EtatPriseDBEnum, nullable=True)
+    )
+    etat_prise_type_chademo: Optional[EtatPriseEnum] = Field(
+        sa_column=SAColumn(EtatPriseDBEnum, nullable=True)
+    )
+    etat_prise_type_ef: Optional[EtatPriseEnum] = Field(
+        sa_column=SAColumn(EtatPriseDBEnum, nullable=True)
+    )
+
     # Relationships
     point_de_charge_id: Optional[UUID] = Field(
         default=None, foreign_key="pointdecharge.id"
diff --git a/src/api/qualicharge/schemas/sql.py b/src/api/qualicharge/schemas/sql.py
