feat: add SQLModels for FlatAnwendungshandbuch; Add PoC for SQLite (#557)

---------

Co-authored-by: Konstantin <[email protected]>

Author: hf-kklein

pyproject.toml

@@ -49,14 +49,20 @@ dynamic = ["readme", "version"]
 kohlrahbi = "kohlrahbi:cli"
 
 [project.optional-dependencies]
+sqlmodels = [
+  "sqlmodel>=0.0.22",
+  "sqlalchemy[mypy]>=2.0.37"
+]
 dev = [
+  "kohlrahbi[sqlmodels]",
   "kohlrahbi[test]",
   "kohlrahbi[lint]",
   "kohlrahbi[typecheck]",
   "kohlrahbi[formatting]",
 ]
 lint = ["pylint==3.3.3", "pylint-pydantic==0.3.5"]
 test = [
+  "kohlrahbi[sqlmodels]",
   "coverage==7.6.10",
   "dictdiffer==0.9.0",
   "freezegun==1.5.1",
@@ -65,6 +71,7 @@ test = [
   "syrupy==4.8.1",
 ]
 typecheck = [
+  "kohlrahbi[sqlmodels]",
   "mypy==1.11.2",
   "networkx-stubs==0.0.1",
   "pandas-stubs==2.2.3.241126",

src/kohlrahbi/models/sqlmodels/__init__.py

@@ -0,0 +1,4 @@
+"""
+sql models to store Anwendungshandbuecher in a database
+you need to install kohlrahbi[sqlmodels] to use this sub-package
+"""

src/kohlrahbi/models/sqlmodels/anwendungshandbuch.py

@@ -0,0 +1,90 @@
+"""model classes copy-pasted from kohlrahbi"""
+
+# pylint: disable=too-few-public-methods
+
+try:
+    from sqlalchemy import UniqueConstraint
+    from sqlmodel import Field, Relationship, SQLModel
+except ImportError as import_error:
+    import_error.msg += "; Did you install kohlrahbi[sqlmodels]?"
+    # sqlmodel is only an optional dependency when kohlrahbi is used to fill a database
+    raise
+import uuid
+from typing import Optional
+from uuid import UUID
+
+from efoli import EdifactFormat, EdifactFormatVersion
+
+# the models here do NOT inherit from the original models, because I didn't manage to fix the issue that arise:
+# https://github.com/Hochfrequenz/kohlrahbi/issues/556
+# <frozen abc>:106: in __new__
+#     ???
+# E   TypeError: FlatAnwendungshandbuch.__init_subclass__() takes no keyword arguments
+#
+# or
+#
+# ..\.tox\dev\Lib\site-packages\sqlmodel\main.py:697: in get_sqlalchemy_type
+#     raise ValueError(f"{type_} has no matching SQLAlchemy type")
+# E   ValueError: <class 'kohlrahbi.models.anwendungshandbuch.AhbMetaInformation'> has no matching SQLAlchemy type
+
+
+class FlatAnwendungshandbuch(SQLModel, table=True):  # team
+    """
+    A flat Anwendungshandbuch (AHB) models an Anwendungshandbuch as combination of some meta data and an ordered list
+    of `.class:`.FlatAhbLine s. Basically a flat Anwendungshandbuch is the result of a simple scraping approach.
+    You can create instances of this class without knowing anything about the "under the hood" structure of AHBs or MIGs
+    """
+
+    id: UUID = Field(primary_key=True, default_factory=uuid.uuid4, description="optional key")
+    meta: Optional["AhbMetaInformation"] = Relationship(back_populates="flatanwendungshandbuch")
+    lines: list["AhbLine"] = Relationship(back_populates="flatanwendungshandbuch")
+
+
+class AhbLine(SQLModel, table=True):
+    """
+    An AhbLine is a single line inside the machine-readable, flat AHB.
+    """
+
+    __table_args__ = (UniqueConstraint("ahb_id", "position_inside_ahb", name="IX_position_once_per_ahb"),)
+    id: UUID = Field(primary_key=True, default_factory=uuid.uuid4, description="optional key")
+    # yes, it's actually that bad already
+
+    position_inside_ahb: int = Field(index=True)
+    ahb_id: UUID | None = Field(default=None, foreign_key="flatanwendungshandbuch.id")
+    flatanwendungshandbuch: FlatAnwendungshandbuch | None = Relationship(back_populates="lines")
+
+    # fields copy-pasted from original model:
+    guid: Optional[UUID]
+    segment_group_key: Optional[str]
+    segment_code: Optional[str]
+    data_element: Optional[str]
+    segment_id: Optional[str]
+    value_pool_entry: Optional[str]
+    name: Optional[str]
+    ahb_expression: Optional[str]
+    conditions: Optional[str]
+    section_name: Optional[str]
+    index: Optional[int]
+
+
+class AhbMetaInformation(SQLModel, table=True):
+    """
+    Meta information about an AHB like e.g. its title, Prüfidentifikator, possible sender and receiver roles
+    """
+
+    __table_args__ = (
+        UniqueConstraint("pruefidentifikator", "edifact_format_version", name="IX_pruefi_once_per_format_version"),
+    )
+    id: UUID = Field(primary_key=True, default_factory=uuid.uuid4, description="optional key")
+    edifact_format: EdifactFormat = Field(index=True)
+    edifact_format_version: EdifactFormatVersion = Field(index=True)
+    flatanwendungshandbuch: Optional[FlatAnwendungshandbuch] = Relationship(
+        back_populates="meta", sa_relationship_kwargs={"uselist": False}
+    )
+    ahb_id: UUID | None = Field(default=None, foreign_key="flatanwendungshandbuch.id")
+
+    # copy-pasted fields from original model:
+    pruefidentifikator: str
+    maus_version: Optional[str]
+    description: Optional[str]
+    direction: Optional[str]

unittests/test_sqlmodels.py

@@ -0,0 +1,83 @@
+"""
+we try to fill a database using kohlrahbi[sqlmodels] and the data from the machine-readable AHB submodule
+"""
+
+import json
+from pathlib import Path
+from typing import Generator
+
+import pytest
+from efoli import EdifactFormat, EdifactFormatVersion
+from sqlmodel import Session, SQLModel, create_engine, select
+
+from kohlrahbi.ahb import process_pruefi
+from kohlrahbi.enums.ahbexportfileformat import AhbExportFileFormat
+from kohlrahbi.models.sqlmodels.anwendungshandbuch import AhbLine, AhbMetaInformation, FlatAnwendungshandbuch
+
+
+@pytest.fixture()
+def sqlite_session(tmp_path: Path) -> Generator[Session, None, None]:
+    database_path = tmp_path / "test.db"
+    engine = create_engine(f"sqlite:///{database_path}")
+    SQLModel.metadata.drop_all(engine)
+    SQLModel.metadata.create_all(engine)
+    with Session(bind=engine) as session:
+        yield session
+
+
+def _load_flat_ahb_to_db(
+    session: Session, json_path: Path, edifact_format: EdifactFormat, edifact_format_version: EdifactFormatVersion
+) -> str:
+    """returns the pruefi"""
+    # you may use this kind of function to eventually load _all_ AHBs into a DB
+    # if you e.g. check out our machine-readable AHB repository,
+    #  for json_path in _mr_ahb_submodule_path.rglob("FV*/**/*.json"):
+    #      with open(json_path, "r", encoding="utf-8") as json_file:
+    #          edifact_format = EdifactFormat(json_path.parent.parent.name)
+    #          edifact_format_version = EdifactFormatVersion(json_path.parent.parent.parent.name)
+    #          _load_flat_ahb_to_db(...)
+    with open(json_path, "r", encoding="utf-8") as json_file:
+        file_body = json.loads(json_file.read())
+        flat_ahb = FlatAnwendungshandbuch.model_validate(file_body)
+        edifact_format = EdifactFormat.UTILMD
+        edifact_format_version = EdifactFormatVersion.FV2504
+        file_body["meta"]["edifact_format"] = str(edifact_format)
+        file_body["meta"]["edifact_format_version"] = str(edifact_format_version)
+        meta = AhbMetaInformation.model_validate(file_body["meta"])
+        session.add(meta)
+        flat_ahb.meta = meta
+        for line_index, raw_line in enumerate(file_body["lines"]):
+            raw_line["position_inside_ahb"] = line_index
+            line = AhbLine.model_validate(raw_line)
+            flat_ahb.lines.append(line)
+            session.add(line)
+        session.add(flat_ahb)
+        session.commit()
+    session.flush()
+    return meta.pruefidentifikator
+
+
+def test_sqlmodels(sqlite_session: Session, tmp_path: Path) -> None:
+    flat_ahb_dir_path = tmp_path
+    docx_file_path = (
+        Path(__file__).parent.parent
+        / "edi_energy_mirror"
+        / "edi_energy_de"
+        / "FV2504"
+        / "UTILMDAHBStrom-informatorischeLesefassung2.1_99991231_20250404.docx"
+    )
+    process_pruefi("55001", docx_file_path, flat_ahb_dir_path, (AhbExportFileFormat.FLATAHB,))
+    flat_ahb_path = flat_ahb_dir_path / "UTILMD" / "flatahb" / "55001.json"
+    assert flat_ahb_path.exists()
+    pruefi = _load_flat_ahb_to_db(sqlite_session, flat_ahb_path, EdifactFormat.UTILMD, EdifactFormatVersion.FV2504)
+    ahbs_from_db = sqlite_session.exec(
+        select(FlatAnwendungshandbuch)
+        .join(AhbMetaInformation, FlatAnwendungshandbuch.id == AhbMetaInformation.ahb_id)  # type:ignore[arg-type]
+        .join(AhbLine, FlatAnwendungshandbuch.id == AhbLine.ahb_id)  # type:ignore[arg-type]
+        .where(AhbMetaInformation.pruefidentifikator == "55001")
+    ).all()
+    ahb_from_db: FlatAnwendungshandbuch = ahbs_from_db[0]
+    assert isinstance(ahb_from_db, FlatAnwendungshandbuch)
+    assert any(ahb_from_db.lines)
+    assert ahb_from_db.meta is not None
+    assert ahb_from_db.meta.pruefidentifikator == pruefi