debt(deps): sqlalchemy 2.0.* upgrade, second attempt

.github/workflows/python.yml

@@ -8,7 +8,7 @@ jobs:
       # Minimum code coverage per file
       COVERAGE_SINGLE: 50
       # Minimum total code coverage
-      COVERAGE_TOTAL: 55
+      COVERAGE_TOTAL: 50
     runs-on: ubuntu-latest
     services:
       postgres:

data/dispatch-sample-data.dump

@@ -3,7 +3,7 @@
 --
 
 -- Dumped from database version 14.6 (Debian 14.6-1.pgdg110+1)
--- Dumped by pg_dump version 14.13 (Homebrew)
+-- Dumped by pg_dump version 14.17 (Homebrew)
 
 SET statement_timeout = 0;
 SET lock_timeout = 0;
@@ -2420,7 +2420,8 @@ CREATE TABLE dispatch_organization_default."case" (
     signal_thread_ts character varying,
     reporter_id integer,
     dedicated_channel boolean,
-    genai_analysis jsonb DEFAULT '{}'::jsonb NOT NULL
+    genai_analysis jsonb DEFAULT '{}'::jsonb NOT NULL,
+    event boolean
 );
 
 
@@ -3089,6 +3090,7 @@ CREATE TABLE dispatch_organization_default.event (
     search_vector tsvector,
     updated_at timestamp without time zone,
     created_at timestamp without time zone,
+    signal_id integer,
     dispatch_user_id integer,
     case_id integer,
     type character varying,
@@ -3325,6 +3327,7 @@ CREATE TABLE dispatch_organization_default.incident (
     reporter_id integer,
     liaison_id integer,
     scribe_id integer,
+    summary character varying,
     incident_document_id integer,
     incident_review_document_id integer,
     tactical_group_id integer,
@@ -3450,6 +3453,7 @@ CREATE TABLE dispatch_organization_default.incident_priority (
     name character varying,
     description character varying,
     page_commander boolean,
+    disable_delayed_message_warning boolean,
     tactical_report_reminder integer DEFAULT 24,
     executive_report_reminder integer DEFAULT 24,
     enabled boolean,
@@ -3619,6 +3623,8 @@ CREATE TABLE dispatch_organization_default.incident_type (
     "default" boolean,
     visibility character varying,
     plugin_metadata json,
+    exclude_from_reminders boolean,
+    exclude_from_review boolean,
     incident_template_document_id integer,
     executive_template_document_id integer,
     review_template_document_id integer,
@@ -4004,6 +4010,7 @@ ALTER SEQUENCE dispatch_organization_default.plugin_instance_id_seq OWNED BY dis
 CREATE TABLE dispatch_organization_default.project (
     id integer NOT NULL,
     name character varying,
+    display_name character varying,
     description character varying,
     "default" boolean,
     color character varying,
@@ -4026,7 +4033,8 @@ CREATE TABLE dispatch_organization_default.project (
     weekly_report_notification_id integer,
     report_incident_instructions character varying,
     report_incident_title_hint character varying,
-    report_incident_description_hint character varying
+    report_incident_description_hint character varying,
+    snooze_extension_oncall_service_id integer
 );
 
 
@@ -4267,6 +4275,7 @@ CREATE TABLE dispatch_organization_default.service (
     updated_at timestamp without time zone,
     created_at timestamp without time zone,
     evergreen boolean,
+    shift_hours_type integer,
     evergreen_owner character varying,
     evergreen_reminder_interval integer,
     evergreen_last_reminder_at timestamp without time zone,
@@ -4549,8 +4558,10 @@ CREATE TABLE dispatch_organization_default.signal_instance (
     created_at timestamp without time zone,
     filter_action character varying,
     engagement_thread_ts character varying,
+    conversation_target character varying,
     case_type_id integer,
     case_priority_id integer,
+    oncall_service_id integer,
     canary boolean
 );
 
@@ -7665,7 +7676,7 @@ COPY dispatch_core.plugin_event (id, name, slug, description, plugin_id, search_
 --
 
 COPY dispatch_organization_default.alembic_version (version_num) FROM stdin;
-928b725d64f6
+8f324b0f365a
 \.
 
 
@@ -13117,6 +13128,12 @@ ALTER TABLE ONLY dispatch_organization_default.event
 ALTER TABLE ONLY dispatch_organization_default.event
     ADD CONSTRAINT event_individual_id_fkey FOREIGN KEY (individual_id) REFERENCES dispatch_organization_default.individual_contact(id) ON DELETE CASCADE;
 
+--
+-- Name: event event_signal_id_fkey; Type: FK CONSTRAINT; Schema: dispatch_organization_default; Owner: postgres
+--
+
+ALTER TABLE ONLY dispatch_organization_default.event
+    ADD CONSTRAINT event_signal_id_fkey FOREIGN KEY (signal_id) REFERENCES dispatch_organization_default.signal(id) ON DELETE CASCADE;
 
 --
 -- Name: feedback feedback_incident_id_fkey; Type: FK CONSTRAINT; Schema: dispatch_organization_default; Owner: postgres
@@ -13853,6 +13870,12 @@ ALTER TABLE ONLY dispatch_organization_default.signal_instance
 ALTER TABLE ONLY dispatch_organization_default.signal_instance
     ADD CONSTRAINT signal_instance_signal_id_fkey FOREIGN KEY (signal_id) REFERENCES dispatch_organization_default.signal(id);
 
+--
+-- Name: signal_instance signal_instance_oncall_service_id_fkey; Type: FK CONSTRAINT; Schema: dispatch_organization_default; Owner: postgres
+--
+
+ALTER TABLE ONLY dispatch_organization_default.signal_instance
+    ADD CONSTRAINT signal_instance_oncall_service_id_fkey FOREIGN KEY (oncall_service_id) REFERENCES dispatch_organization_default.service(id);
 
 --
 -- Name: signal signal_oncall_service_id_fkey; Type: FK CONSTRAINT; Schema: dispatch_organization_default; Owner: postgres
@@ -15009,4 +15032,3 @@ ALTER TABLE ONLY public.workflow_term
 --
 -- PostgreSQL database dump complete
 --
-

requirements-base.in

@@ -53,7 +53,7 @@ slowapi
 spacy
 sqlalchemy-filters
 sqlalchemy-utils
-sqlalchemy<1.4  # NOTE temporarily until https://github.com/kvesteri/sqlalchemy-utils/issues/505 is fixed
+sqlalchemy==2.0.8
 statsmodels
 tabulate
 tenacity

requirements-base.txt

@@ -5,7 +5,6 @@
 #    pip-compile --output-file=requirements-base.txt requirements-base.in
 #
 --index-url https://pypi.netflix.net/simple
---trusted-host pypi.org
 
 aiocache==0.12.3
     # via -r requirements-base.in
@@ -449,7 +448,7 @@ spacy-legacy==3.0.12
     # via spacy
 spacy-loggers==1.0.5
     # via spacy
-sqlalchemy==1.3.24
+sqlalchemy==2.0.8
     # via
     #   -r requirements-base.in
     #   alembic
@@ -506,6 +505,7 @@ typing-extensions==4.13.2
     #   openai
     #   pydantic
     #   schemathesis
+    #   sqlalchemy
     #   typer
 tzdata==2025.1
     # via pandas

src/dispatch/auth/models.py

@@ -115,7 +115,7 @@ class DispatchUserProject(Base, TimeStampMixin):
     dispatch_user = relationship(DispatchUser, backref="projects")
 
     project_id = Column(Integer, ForeignKey(Project.id), primary_key=True)
-    project = relationship(Project, backref="users")
+    project = relationship(Project, backref="users", overlaps="dispatch_user_project")
 
     default = Column(Boolean, default=False)
 

src/dispatch/database/core.py

@@ -15,8 +15,7 @@
 from pydantic.error_wrappers import ErrorWrapper, ValidationError
 from sqlalchemy import create_engine, inspect
 from sqlalchemy.engine.url import make_url
-from sqlalchemy.ext.declarative import declarative_base, declared_attr
-from sqlalchemy.orm import Session, object_session, sessionmaker
+from sqlalchemy.orm import Session, object_session, sessionmaker, DeclarativeBase, declared_attr
 from sqlalchemy.sql.expression import true
 from sqlalchemy_utils import get_mapper
 from starlette.requests import Request
@@ -96,13 +95,14 @@ def resolve_attr(obj, attr, default=None):
         return default
 
 
-class CustomBase:
+class Base(DeclarativeBase):
+    """Base class for all SQLAlchemy models."""
     __repr_attrs__ = []
     __repr_max_length__ = 15
 
-    @declared_attr
-    def __tablename__(self):
-        return resolve_table_name(self.__name__)
+    @declared_attr.directive
+    def __tablename__(cls):
+        return resolve_table_name(cls.__name__)
 
     def dict(self):
         """Returns a dict representation of a model."""
@@ -149,9 +149,6 @@ def __repr__(self):
             id_str,
             " " + self._repr_attrs_str if self._repr_attrs_str else "",
         )
-
-
-Base = declarative_base(cls=CustomBase)
 make_searchable(Base.metadata)
 
 
@@ -177,10 +174,11 @@ def get_class_by_tablename(table_fullname: str) -> Any:
     """Return class reference mapped to table."""
 
     def _find_class(name):
-        for c in Base._decl_class_registry.values():
-            if hasattr(c, "__table__"):
-                if c.__table__.fullname.lower() == name.lower():
-                    return c
+        for mapper in Base.registry.mappers:
+            cls = mapper.class_
+            if hasattr(cls, "__table__"):
+                if cls.__table__.fullname.lower() == name.lower():
+                    return cls
 
     mapped_name = resolve_table_name(table_fullname)
     mapped_class = _find_class(mapped_name)

src/dispatch/database/manage.py

@@ -4,8 +4,9 @@
 from alembic import command as alembic_command
 from alembic.config import Config as AlembicConfig
 
-from sqlalchemy import text
-from sqlalchemy.schema import CreateSchema
+from sqlalchemy import Engine, text
+from sqlalchemy.engine import Connection
+from sqlalchemy.schema import CreateSchema, Table
 from sqlalchemy_utils import create_database, database_exists
 
 from dispatch import config
@@ -33,40 +34,40 @@ def version_schema(script_location: str):
     alembic_command.stamp(alembic_cfg, "head")
 
 
-def get_core_tables():
+def get_core_tables()  -> list[Table]:
     """Fetches tables that belong to the 'dispatch_core' schema."""
-    core_tables = []
+    core_tables: list[Table] = []
     for _, table in Base.metadata.tables.items():
         if table.schema == "dispatch_core":
             core_tables.append(table)
     return core_tables
 
 
-def get_tenant_tables():
+def get_tenant_tables() -> list[Table]:
     """Fetches tables that belong to their own tenant tables."""
-    tenant_tables = []
+    tenant_tables: list[Table] = []
     for _, table in Base.metadata.tables.items():
         if not table.schema:
             tenant_tables.append(table)
     return tenant_tables
 
 
-def init_database(engine):
+def init_database(engine: Engine):
     """Initializes the database."""
     if not database_exists(str(config.SQLALCHEMY_DATABASE_URI)):
         create_database(str(config.SQLALCHEMY_DATABASE_URI))
 
     schema_name = "dispatch_core"
-    if not engine.dialect.has_schema(engine, schema_name):
-        with engine.connect() as connection:
-            connection.execute(CreateSchema(schema_name))
+    with engine.begin() as connection:
+        connection.execute(CreateSchema(schema_name, if_not_exists=True))
 
     tables = get_core_tables()
 
     Base.metadata.create_all(engine, tables=tables)
 
     version_schema(script_location=config.ALEMBIC_CORE_REVISION_PATH)
-    setup_fulltext_search(engine, tables)
+    with engine.connect() as connection:
+        setup_fulltext_search(connection, tables)
 
     # setup an required database functions
     session = sessionmaker(bind=engine)
@@ -133,24 +134,21 @@ def init_database(engine):
         )
 
 
-def init_schema(*, engine, organization: Organization):
+def init_schema(*, engine: Engine, organization: Organization) -> Organization:
     """Initializes a new schema."""
     schema_name = f"{DISPATCH_ORGANIZATION_SCHEMA_PREFIX}_{organization.slug}"
 
-    if not engine.dialect.has_schema(engine, schema_name):
-        with engine.connect() as connection:
-            connection.execute(CreateSchema(schema_name))
+    with engine.begin() as connection:
+        connection.execute(CreateSchema(schema_name, if_not_exists=True))
 
     # set the schema for table creation
     tables = get_tenant_tables()
 
-    schema_engine = engine.execution_options(
-        schema_translate_map={
-            None: schema_name,
-        }
-    )
+    # alter each table's schema
+    for t in tables:
+        t.schema = schema_name
 
-    Base.metadata.create_all(schema_engine, tables=tables)
+    Base.metadata.create_all(engine, tables=tables)
 
     # put schema under version control
     version_schema(script_location=config.ALEMBIC_TENANT_REVISION_PATH)
@@ -163,7 +161,7 @@ def init_schema(*, engine, organization: Organization):
 
         setup_fulltext_search(connection, tables)
 
-    session = sessionmaker(bind=schema_engine)
+    session = sessionmaker(bind=engine)
     db_session = session()
 
     organization = db_session.merge(organization)
@@ -172,7 +170,7 @@ def init_schema(*, engine, organization: Organization):
     return organization
 
 
-def setup_fulltext_search(connection, tables):
+def setup_fulltext_search(connection: Connection, tables: list[Table]) -> None:
     """Syncs any required fulltext table triggers and functions."""
     # parsing functions
     function_path = os.path.join(

src/dispatch/database/revisions/core/env.py

@@ -1,5 +1,5 @@
 from alembic import context
-from sqlalchemy import engine_from_config, pool
+from sqlalchemy import engine_from_config, pool, text
 
 from dispatch.logging import logging
 from dispatch.config import SQLALCHEMY_DATABASE_URI
@@ -51,7 +51,8 @@ def process_revision_directives(context, revision, directives):
     log.info("Migrating dispatch core schema...")
     # migrate common tables
     with connectable.connect() as connection:
-        connection.execute(f'set search_path to "{CORE_SCHEMA_NAME}"')
+        set_search_path = text(f'set search_path to "{CORE_SCHEMA_NAME}"')
+        connection.execute(set_search_path)
         connection.dialect.default_schema_name = CORE_SCHEMA_NAME
         context.configure(
             connection=connection,

src/dispatch/database/revisions/tenant/env.py

@@ -1,5 +1,5 @@
 from alembic import context
-from sqlalchemy import engine_from_config, pool, inspect
+from sqlalchemy import engine_from_config, pool, inspect, text
 
 
 from dispatch.logging import logging
@@ -58,7 +58,8 @@ def process_revision_directives(context, revision, directives):
         # get the schema names
         for schema in get_tenant_schemas(connection):
             log.info(f"Migrating {schema}...")
-            connection.execute(f'set search_path to "{schema}"')
+            set_search_path = text(f'set search_path to "{schema}"')
+            connection.execute(set_search_path)
             connection.dialect.default_schema_name = schema
 
             context.configure(

src/dispatch/database/revisions/tenant/versions/2023-03-03_7ddae3ba7822.py

@@ -20,8 +20,9 @@
 
 def upgrade():
     conn = op.get_context().connection
-    metadata = MetaData(bind=conn, schema=conn.dialect.default_schema_name)
-    table = sa.Table("signal", metadata, autoload=True)
+    metadata = MetaData(schema=conn.dialect.default_schema_name)
+    metadata.bind = conn
+    table = sa.Table("signal", metadata, autoload_with=conn)
     sync_trigger(conn, table, "search_vector", ["name", "description", "variant"])