Migrate from declarative_base to DeclarativeBase and improve typing

Co-authored-by: chrisburr <5220533+chrisburr@users.noreply.github.com>

Author: Copilot

diracx-db/src/diracx/db/sql/auth/schema.py

@@ -8,13 +8,15 @@
     String,
     Uuid,
 )
-from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm import DeclarativeBase
 
 from diracx.db.sql.utils import Column, DateNowColumn, EnumColumn, NullColumn
 
 USER_CODE_LENGTH = 8
 
-Base = declarative_base()
+
+class Base(DeclarativeBase):
+    pass
 
 
 class FlowStatus(Enum):

diracx-db/src/diracx/db/sql/dummy/schema.py

@@ -3,11 +3,13 @@
 from __future__ import annotations
 
 from sqlalchemy import ForeignKey, Integer, String, Uuid
-from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm import DeclarativeBase
 
 from diracx.db.sql.utils import Column, DateNowColumn
 
-Base = declarative_base()
+
+class Base(DeclarativeBase):
+    pass
 
 
 class Owners(Base):

diracx-db/src/diracx/db/sql/job/schema.py

@@ -9,11 +9,13 @@
     String,
     Text,
 )
-from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm import DeclarativeBase
 
 from ..utils import Column, EnumBackedBool, NullColumn
 
-JobDBBase = declarative_base()
+
+class JobDBBase(DeclarativeBase):
+    pass
 
 
 class AccountedFlagEnum(types.TypeDecorator):

diracx-db/src/diracx/db/sql/job_logging/schema.py

@@ -3,11 +3,13 @@
 from datetime import UTC, datetime
 
 from sqlalchemy import Integer, Numeric, PrimaryKeyConstraint, String, TypeDecorator
-from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm import DeclarativeBase
 
 from ..utils import Column, DateNowColumn
 
-JobLoggingDBBase = declarative_base()
+
+class JobLoggingDBBase(DeclarativeBase):
+    pass
 
 
 class MagicEpochDateTime(TypeDecorator):

diracx-db/src/diracx/db/sql/pilot_agents/schema.py

@@ -8,11 +8,13 @@
     String,
     Text,
 )
-from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm import DeclarativeBase
 
 from ..utils import Column, EnumBackedBool, NullColumn
 
-PilotAgentsDBBase = declarative_base()
+
+class PilotAgentsDBBase(DeclarativeBase):
+    pass
 
 
 class PilotAgents(PilotAgentsDBBase):

diracx-db/src/diracx/db/sql/sandbox_metadata/schema.py

@@ -9,11 +9,13 @@
     String,
     UniqueConstraint,
 )
-from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm import DeclarativeBase
 
 from diracx.db.sql.utils import Column, DateNowColumn
 
-Base = declarative_base()
+
+class Base(DeclarativeBase):
+    pass
 
 
 class SBOwners(Base):

diracx-db/src/diracx/db/sql/task_queue/schema.py

@@ -9,11 +9,13 @@
     Integer,
     String,
 )
-from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm import DeclarativeBase
 
 from ..utils import Column
 
-TaskQueueDBBase = declarative_base()
+
+class TaskQueueDBBase(DeclarativeBase):
+    pass
 
 
 class TaskQueues(TaskQueueDBBase):

diracx-db/src/diracx/db/sql/utils/base.py

@@ -8,13 +8,16 @@
 from collections.abc import AsyncIterator
 from contextvars import ContextVar
 from datetime import datetime
-from typing import Any, Self, cast
+from typing import Any, Self, cast, TYPE_CHECKING
 
 from pydantic import TypeAdapter
 from sqlalchemy import DateTime, MetaData, func, select
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
 
+if TYPE_CHECKING:
+    from sqlalchemy.orm import DeclarativeBase
+
 from diracx.core.exceptions import InvalidQueryError
 from diracx.core.extensions import select_from_extension
 from diracx.core.models import (
@@ -233,7 +236,7 @@ async def ping(self):
 
     async def _search(
         self,
-        table: Any,
+        table: type[DeclarativeBase],
         parameters: list[str] | None,
         search: list[SearchSpec],
         sorts: list[SortSpec],
@@ -273,7 +276,7 @@ async def _search(
         ]
 
     async def _summary(
-        self, table: Any, group_by: list[str], search: list[SearchSpec]
+        self, table: type[DeclarativeBase], group_by: list[str], search: list[SearchSpec]
     ) -> list[dict[str, str | int]]:
         """Get a summary of the elements of a table."""
         columns = _get_columns(table.__table__, group_by)

extensions/gubbins/gubbins-db/src/gubbins/db/sql/jobs/schema.py

@@ -1,4 +1,4 @@
-from diracx.db.sql.job.db import JobDBBase
+from diracx.db.sql.job.schema import JobDBBase
 from diracx.db.sql.utils import Column
 from sqlalchemy import (
     ForeignKey,

extensions/gubbins/gubbins-db/src/gubbins/db/sql/lollygag/schema.py

@@ -2,9 +2,11 @@
 # in place of the SQLAlchemy one. Have a look at them
 from diracx.db.sql.utils import Column, DateNowColumn
 from sqlalchemy import ForeignKey, Integer, String, Uuid
-from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm import DeclarativeBase
 
-Base = declarative_base()
+
+class Base(DeclarativeBase):
+    pass
 
 
 class Owners(Base):