diff --git a/src/api/qualicharge/migrations/versions/3e8ea07e5b66_update_db_enum_values.py b/src/api/qualicharge/migrations/versions/3e8ea07e5b66_update_db_enum_values.py new file mode 100644 index 00000000..df8386b0 --- /dev/null +++ b/src/api/qualicharge/migrations/versions/3e8ea07e5b66_update_db_enum_values.py @@ -0,0 +1,264 @@ +"""Update db enum values + +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 Session, SQLModel, update + +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", + ) + + # Update data + with Session(connection) as session: + for key, value in enum_to_dict(enum_, revert).items(): + session.exec( + update(schema) + .where(getattr(schema, column_name) == key) + .values({column_name: value}) + ) + + # Alter table column to a the new ENUM + op.alter_column( + schema.__tablename__, + column_name, + existing_type=VARCHAR, + type_=postgresql.ENUM( + *enum_to_dict(enum_, revert).values(), + name=new_enum_db_name, + ), + existing_nullable=False, + postgresql_using=f"{column_name}::{new_enum_db_name}", + ) + + # Delete old ENUM database type + postgresql.ENUM( + *enum_to_dict(enum_, not revert).values(), + name=existing_enum_db_name, + ).drop(connection, checkfirst=True) + + +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 index ff15876e..34f5c9eb 100644 --- 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,24 @@ 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) + ) + horodatage: PastDatetime + 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 index c528d35e..e81292ae 100644 --- a/src/api/qualicharge/schemas/sql.py +++ b/src/api/qualicharge/schemas/sql.py @@ -21,15 +21,11 @@ from ..models.static import Statique from . import BaseTimestampedSQLModel from .core import ( - AccessibilitePMREnum, Amenageur, - ConditionAccesEnum, Enseigne, - ImplantationStationEnum, Localisation, Operateur, PointDeCharge, - RaccordementEnum, Station, ) @@ -43,7 +39,7 @@ def __init__(self, df: pd.DataFrame, connection: Connection): """Add table cache keys.""" logger.info("Loading input dataframe containing %d rows", len(df)) - self._statique: pd.DataFrame = self._fix_enums(df) + self._statique: pd.DataFrame = df self._statique_with_fk: pd.DataFrame = self._statique.copy() self._saved_schemas: list[type[BaseTimestampedSQLModel]] = [] @@ -84,24 +80,6 @@ def _get_schema_fks(schema: type[BaseTimestampedSQLModel]) -> list[str]: for fk in schema.metadata.tables[schema.__tablename__].foreign_keys # type: ignore[index] ] - def _fix_enums(self, df: pd.DataFrame) -> pd.DataFrame: - """Fix enums representation in dataframe.""" - logger.debug("Fixing enum columns representation") - target = [] - src = [] - - for enum_ in ( - ImplantationStationEnum, - ConditionAccesEnum, - AccessibilitePMREnum, - RaccordementEnum, - ): - for entry in enum_: - target.append(str(entry.name)) - src.append(entry.value) - - return df.replace(to_replace=src, value=target) - def _get_fields_for_schema( self, schema: type[BaseTimestampedSQLModel], with_fk: bool = False ) -> list[str]: