Add custom sqlacodegen generator

Author: gfrn

.github/workflows/update-orm.yml

@@ -100,7 +100,7 @@ jobs:
         if: ${{ steps.checkUpdate.outputs.needsUpdate == 'true' }}
         run: |
           set -eux
-          sqlacodegen mysql+mysqlconnector://root:[email protected]/ispyb_build --options nojoined --outfile _auto_db_schema.py.in
+          sqlacodegen --generator declarative_ispyb mysql+mysqlconnector://root:[email protected]/ispyb_build --options nojoined --outfile _auto_db_schema.py.in
           # This code produces false positives due to non-deterministic ordering.
           # Add an identifier to the file, and only update/PR the changes when the
           # identifier indicates the file is out of date.

pyproject.toml

@@ -36,6 +36,9 @@ Bug-Tracker = "https://github.com/DiamondLightSource/ispyb-api/issues"
 "ispyb.job" = "ispyb.job"
 "ispyb.last_data_collections_on" = "ispyb.last_data_collections_on"
 
+[project.entry-points."sqlacodegen.generators"]
+declarative_ispyb = "ispyb.generators:DeclarativeIspybGenerator"
+
 [project.entry-points."libtbx.precommit"]
 ispyb = "ispyb"
 

src/ispyb/generators.py

@@ -0,0 +1,58 @@
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import inflect
+from sqlacodegen.generators import DeclarativeGenerator
+from sqlacodegen.models import RelationshipAttribute, RelationshipType
+from sqlalchemy import (
+    MetaData,
+)
+from sqlalchemy.engine import Connection, Engine
+
+
+class DeclarativeIspybGenerator(DeclarativeGenerator):
+    def __init__(
+        self,
+        metadata: MetaData,
+        bind: Connection | Engine,
+        options: Sequence[str],
+        *,
+        indentation: str = "    ",
+        base_class_name: str = "Base",
+    ):
+        super().__init__(metadata, bind, options, indentation=indentation)
+        self.base_class_name: str = base_class_name
+        self.inflect_engine = inflect.engine()
+
+    def generate_relationship_name(
+        self,
+        relationship: RelationshipAttribute,
+        global_names: set[str],
+        local_names: set[str],
+    ) -> None:
+        # Self referential reverse relationships
+        preferred_name: str
+        if (
+            relationship.type
+            in (RelationshipType.ONE_TO_MANY, RelationshipType.ONE_TO_ONE)
+            and relationship.source is relationship.target
+            and relationship.backref
+            and relationship.backref.name
+        ):
+            preferred_name = relationship.backref.name + "_reverse"
+        else:
+            preferred_name = relationship.target.table.name
+
+            if relationship.constraint:
+                is_source = relationship.source.table is relationship.constraint.table
+                if is_source or relationship.type not in (
+                    RelationshipType.ONE_TO_ONE,
+                    RelationshipType.ONE_TO_MANY,
+                ):
+                    column_names = [c.name for c in relationship.constraint.columns]
+                    if len(column_names) == 1 and column_names[0].endswith("_id"):
+                        preferred_name = column_names[0][:-3]
+
+        relationship.name = preferred_name

src/ispyb/sqlalchemy/__init__.py

@@ -11,24 +11,26 @@
 from ._auto_db_schema import *  # noqa F403; lgtm
 from ._auto_db_schema import (
     AutoProcProgram,
-    AutoProcScaling,
     ProcessingJob,
     __schema_version__,
 )
 
 logger = logging.getLogger("ispyb.sqlalchemy")
 
 AutoProcProgram.AutoProcProgramAttachments = relationship(
-    "AutoProcProgramAttachment", back_populates="AutoProcProgram"
-)
-AutoProcScaling.AutoProcScalingStatistics = relationship(
-    "AutoProcScalingStatistics", back_populates="AutoProcScaling"
+    "AutoProcProgramAttachment",
+    back_populates="AutoProcProgram",
+    overlaps="AutoProcProgramAttachment",
 )
 ProcessingJob.ProcessingJobParameters = relationship(
-    "ProcessingJobParameter", back_populates="ProcessingJob"
+    "ProcessingJobParameter",
+    back_populates="ProcessingJob",
+    overlaps="ProcessingJobParameter",
 )
 ProcessingJob.ProcessingJobImageSweeps = relationship(
-    "ProcessingJobImageSweep", back_populates="ProcessingJob"
+    "ProcessingJobImageSweep",
+    back_populates="ProcessingJob",
+    overlaps="ProcessingJobImageSweep",
 )
 assert __schema_version__