Upgrade SQLAlchemy to 2.0.2

homeassistant/components/logbook/models.py

@@ -82,7 +82,7 @@ class EventAsRow:
 
 
 @callback
-def async_event_to_row(event: Event) -> EventAsRow | None:
+def async_event_to_row(event: Event) -> EventAsRow:
     """Convert an event to a row."""
     if event.event_type != EVENT_STATE_CHANGED:
         return EventAsRow(

homeassistant/components/logbook/processor.py

@@ -1,14 +1,14 @@
 """Event parser and human readable log generator."""
 from __future__ import annotations
 
-from collections.abc import Callable, Generator
+from collections.abc import Callable, Generator, Sequence
 from contextlib import suppress
 from dataclasses import dataclass
 from datetime import datetime as dt
 from typing import Any
 
+from sqlalchemy.engine import Result
 from sqlalchemy.engine.row import Row
-from sqlalchemy.orm.query import Query
 
 from homeassistant.components.recorder.filters import Filters
 from homeassistant.components.recorder.models import (
@@ -70,7 +70,7 @@ class LogbookRun:
     event_cache: EventCache
     entity_name_cache: EntityNameCache
     include_entity_name: bool
-    format_time: Callable[[Row], Any]
+    format_time: Callable[[Row | EventAsRow], Any]
 
 
 class EventProcessor:
@@ -133,13 +133,13 @@ def get_events(
     ) -> list[dict[str, Any]]:
         """Get events for a period of time."""
 
-        def yield_rows(query: Query) -> Generator[Row, None, None]:
+        def yield_rows(result: Result) -> Sequence[Row] | Result:
             """Yield rows from the database."""
             # end_day - start_day intentionally checks .days and not .total_seconds()
             # since we don't want to switch over to buffered if they go
             # over one day by a few hours since the UI makes it so easy to do that.
             if self.limited_select or (end_day - start_day).days <= 1:
-                return query.all()  # type: ignore[no-any-return]
+                return result.all()
             # Only buffer rows to reduce memory pressure
             # if we expect the result set is going to be very large.
             # What is considered very large is going to differ
@@ -149,7 +149,7 @@ def yield_rows(query: Query) -> Generator[Row, None, None]:
             # even and RPi3 that number seems higher in testing
             # so we don't switch over until we request > 1 day+ of data.
             #
-            return query.yield_per(1024)  # type: ignore[no-any-return]
+            return result.yield_per(1024)
 
         stmt = statement_for_request(
             start_day,
@@ -164,12 +164,12 @@ def yield_rows(query: Query) -> Generator[Row, None, None]:
             return self.humanify(yield_rows(session.execute(stmt)))
 
     def humanify(
-        self, row_generator: Generator[Row | EventAsRow, None, None]
+        self, rows: Generator[EventAsRow, None, None] | Sequence[Row] | Result
     ) -> list[dict[str, str]]:
         """Humanify rows."""
         return list(
             _humanify(
-                row_generator,
+                rows,
                 self.ent_reg,
                 self.logbook_run,
                 self.context_augmenter,
@@ -178,7 +178,7 @@ def humanify(
 
 
 def _humanify(
-    rows: Generator[Row | EventAsRow, None, None],
+    rows: Generator[EventAsRow, None, None] | Sequence[Row] | Result,
     ent_reg: er.EntityRegistry,
     logbook_run: LogbookRun,
     context_augmenter: ContextAugmenter,
@@ -263,7 +263,7 @@ def __init__(self, hass: HomeAssistant) -> None:
         self._memorize_new = True
         self._lookup: dict[str | None, Row | EventAsRow | None] = {None: None}
 
-    def memorize(self, row: Row) -> str | None:
+    def memorize(self, row: Row | EventAsRow) -> str | None:
         """Memorize a context from the database."""
         if self._memorize_new:
             context_id: str = row.context_id
@@ -276,7 +276,7 @@ def clear(self) -> None:
         self._lookup.clear()
         self._memorize_new = False
 
-    def get(self, context_id: str) -> Row | None:
+    def get(self, context_id: str) -> Row | EventAsRow | None:
         """Get the context origin."""
         return self._lookup.get(context_id)
 
@@ -294,7 +294,7 @@ def __init__(self, logbook_run: LogbookRun) -> None:
 
     def _get_context_row(
         self, context_id: str | None, row: Row | EventAsRow
-    ) -> Row | EventAsRow:
+    ) -> Row | EventAsRow | None:
         """Get the context row from the id or row context."""
         if context_id:
             return self.context_lookup.get(context_id)

homeassistant/components/logbook/queries/all.py

@@ -2,9 +2,9 @@
 from __future__ import annotations
 
 from sqlalchemy import lambda_stmt
-from sqlalchemy.orm import Query
-from sqlalchemy.sql.elements import ClauseList
+from sqlalchemy.sql.elements import ColumnElement
 from sqlalchemy.sql.lambdas import StatementLambdaElement
+from sqlalchemy.sql.selectable import Select
 
 from homeassistant.components.recorder.db_schema import (
     LAST_UPDATED_INDEX_TS,
@@ -24,8 +24,8 @@ def all_stmt(
     start_day: float,
     end_day: float,
     event_types: tuple[str, ...],
-    states_entity_filter: ClauseList | None = None,
-    events_entity_filter: ClauseList | None = None,
+    states_entity_filter: ColumnElement | None = None,
+    events_entity_filter: ColumnElement | None = None,
     context_id: str | None = None,
 ) -> StatementLambdaElement:
     """Generate a logbook query for all entities."""
@@ -37,16 +37,29 @@ def all_stmt(
         # are gone from the database remove the
         # _legacy_select_events_context_id()
         stmt += lambda s: s.where(Events.context_id == context_id).union_all(
-            _states_query_for_context_id(start_day, end_day, context_id),
-            legacy_select_events_context_id(start_day, end_day, context_id),
+            _states_query_for_context_id(
+                start_day,
+                end_day,
+                # https://github.com/python/mypy/issues/2608
+                context_id,  # type:ignore[arg-type]
+            ),
+            legacy_select_events_context_id(
+                start_day,
+                end_day,
+                # https://github.com/python/mypy/issues/2608
+                context_id,  # type:ignore[arg-type]
+            ),
         )
     else:
         if events_entity_filter is not None:
             stmt += lambda s: s.where(events_entity_filter)
 
         if states_entity_filter is not None:
             stmt += lambda s: s.union_all(
-                _states_query_for_all(start_day, end_day).where(states_entity_filter)
+                _states_query_for_all(start_day, end_day).where(
+                    # https://github.com/python/mypy/issues/2608
+                    states_entity_filter  # type:ignore[arg-type]
+                )
             )
         else:
             stmt += lambda s: s.union_all(_states_query_for_all(start_day, end_day))
@@ -55,20 +68,20 @@ def all_stmt(
     return stmt
 
 
-def _states_query_for_all(start_day: float, end_day: float) -> Query:
+def _states_query_for_all(start_day: float, end_day: float) -> Select:
     return apply_states_filters(_apply_all_hints(select_states()), start_day, end_day)
 
 
-def _apply_all_hints(query: Query) -> Query:
+def _apply_all_hints(sel: Select) -> Select:
     """Force mysql to use the right index on large selects."""
-    return query.with_hint(
+    return sel.with_hint(
         States, f"FORCE INDEX ({LAST_UPDATED_INDEX_TS})", dialect_name="mysql"
     )
 
 
 def _states_query_for_context_id(
     start_day: float, end_day: float, context_id: str
-) -> Query:
+) -> Select:
     return apply_states_filters(select_states(), start_day, end_day).where(
         States.context_id == context_id
     )

homeassistant/components/logbook/queries/common.py

@@ -5,8 +5,7 @@
 
 import sqlalchemy
 from sqlalchemy import select
-from sqlalchemy.orm import Query
-from sqlalchemy.sql.elements import ClauseList
+from sqlalchemy.sql.elements import BooleanClauseList, ColumnElement
 from sqlalchemy.sql.expression import literal
 from sqlalchemy.sql.selectable import Select
 
@@ -69,7 +68,7 @@
     literal(value=None, type_=sqlalchemy.String).label("old_format_icon"),
 )
 
-EVENT_COLUMNS_FOR_STATE_SELECT = [
+EVENT_COLUMNS_FOR_STATE_SELECT = (
     literal(value=None, type_=sqlalchemy.Text).label("event_id"),
     # We use PSEUDO_EVENT_STATE_CHANGED aka None for
     # state_changed events since it takes up less
@@ -84,7 +83,7 @@
     States.context_user_id.label("context_user_id"),
     States.context_parent_id.label("context_parent_id"),
     literal(value=None, type_=sqlalchemy.Text).label("shared_data"),
-]
+)
 
 EMPTY_STATE_COLUMNS = (
     literal(value=0, type_=sqlalchemy.Integer).label("state_id"),
@@ -103,8 +102,8 @@
 
 # Virtual column to tell logbook if it should avoid processing
 # the event as its only used to link contexts
-CONTEXT_ONLY = literal("1").label("context_only")
-NOT_CONTEXT_ONLY = literal(None).label("context_only")
+CONTEXT_ONLY = literal(value="1", type_=sqlalchemy.String).label("context_only")
+NOT_CONTEXT_ONLY = literal(value=None, type_=sqlalchemy.String).label("context_only")
 
 
 def select_events_context_id_subquery(
@@ -188,15 +187,15 @@ def legacy_select_events_context_id(
     )
 
 
-def apply_states_filters(query: Query, start_day: float, end_day: float) -> Query:
+def apply_states_filters(sel: Select, start_day: float, end_day: float) -> Select:
     """Filter states by time range.
 
     Filters states that do not have an old state or new state (added / removed)
     Filters states that are in a continuous domain with a UOM.
     Filters states that do not have matching last_updated_ts and last_changed_ts.
     """
     return (
-        query.filter(
+        sel.filter(
             (States.last_updated_ts > start_day) & (States.last_updated_ts < end_day)
         )
         .outerjoin(OLD_STATE, (States.old_state_id == OLD_STATE.state_id))
@@ -212,18 +211,18 @@ def apply_states_filters(query: Query, start_day: float, end_day: float) -> Quer
     )
 
 
-def _missing_state_matcher() -> sqlalchemy.and_:
+def _missing_state_matcher() -> ColumnElement[bool]:
     # The below removes state change events that do not have
     # and old_state or the old_state is missing (newly added entities)
     # or the new_state is missing (removed entities)
     return sqlalchemy.and_(
-        OLD_STATE.state_id.isnot(None),
+        OLD_STATE.state_id.is_not(None),
         (States.state != OLD_STATE.state),
-        States.state.isnot(None),
+        States.state.is_not(None),
     )
 
 
-def _not_continuous_entity_matcher() -> sqlalchemy.or_:
+def _not_continuous_entity_matcher() -> ColumnElement[bool]:
     """Match non continuous entities."""
     return sqlalchemy.or_(
         # First exclude domains that may be continuous
@@ -236,7 +235,7 @@ def _not_continuous_entity_matcher() -> sqlalchemy.or_:
     )
 
 
-def _not_possible_continuous_domain_matcher() -> sqlalchemy.and_:
+def _not_possible_continuous_domain_matcher() -> ColumnElement[bool]:
     """Match not continuous domains.
 
     This matches domain that are always considered continuous
@@ -254,7 +253,7 @@ def _not_possible_continuous_domain_matcher() -> sqlalchemy.and_:
     ).self_group()
 
 
-def _conditionally_continuous_domain_matcher() -> sqlalchemy.or_:
+def _conditionally_continuous_domain_matcher() -> ColumnElement[bool]:
     """Match conditionally continuous domains.
 
     This matches domain that are only considered
@@ -268,22 +267,22 @@ def _conditionally_continuous_domain_matcher() -> sqlalchemy.or_:
     ).self_group()
 
 
-def _not_uom_attributes_matcher() -> ClauseList:
+def _not_uom_attributes_matcher() -> BooleanClauseList:
     """Prefilter ATTR_UNIT_OF_MEASUREMENT as its much faster in sql."""
     return ~StateAttributes.shared_attrs.like(
         UNIT_OF_MEASUREMENT_JSON_LIKE
     ) | ~States.attributes.like(UNIT_OF_MEASUREMENT_JSON_LIKE)
 
 
-def apply_states_context_hints(query: Query) -> Query:
+def apply_states_context_hints(sel: Select) -> Select:
     """Force mysql to use the right index on large context_id selects."""
-    return query.with_hint(
+    return sel.with_hint(
         States, f"FORCE INDEX ({STATES_CONTEXT_ID_INDEX})", dialect_name="mysql"
     )
 
 
-def apply_events_context_hints(query: Query) -> Query:
+def apply_events_context_hints(sel: Select) -> Select:
     """Force mysql to use the right index on large context_id selects."""
-    return query.with_hint(
+    return sel.with_hint(
         Events, f"FORCE INDEX ({EVENTS_CONTEXT_ID_INDEX})", dialect_name="mysql"
     )

homeassistant/components/logbook/queries/devices.py

@@ -5,10 +5,9 @@
 
 import sqlalchemy
 from sqlalchemy import lambda_stmt, select
-from sqlalchemy.orm import Query
-from sqlalchemy.sql.elements import ClauseList
+from sqlalchemy.sql.elements import BooleanClauseList
 from sqlalchemy.sql.lambdas import StatementLambdaElement
-from sqlalchemy.sql.selectable import CTE, CompoundSelect
+from sqlalchemy.sql.selectable import CTE, CompoundSelect, Select
 
 from homeassistant.components.recorder.db_schema import (
     DEVICE_ID_IN_EVENT,
@@ -32,7 +31,7 @@ def _select_device_id_context_ids_sub_query(
     end_day: float,
     event_types: tuple[str, ...],
     json_quotable_device_ids: list[str],
-) -> CompoundSelect:
+) -> Select:
     """Generate a subquery to find context ids for multiple devices."""
     inner = select_events_context_id_subquery(start_day, end_day, event_types).where(
         apply_event_device_id_matchers(json_quotable_device_ids)
@@ -41,7 +40,7 @@ def _select_device_id_context_ids_sub_query(
 
 
 def _apply_devices_context_union(
-    query: Query,
+    sel: Select,
     start_day: float,
     end_day: float,
     event_types: tuple[str, ...],
@@ -54,7 +53,7 @@ def _apply_devices_context_union(
         event_types,
         json_quotable_device_ids,
     ).cte()
-    return query.union_all(
+    return sel.union_all(
         apply_events_context_hints(
             select_events_context_only()
             .select_from(devices_cte)
@@ -91,7 +90,7 @@ def devices_stmt(
 
 def apply_event_device_id_matchers(
     json_quotable_device_ids: Iterable[str],
-) -> ClauseList:
+) -> BooleanClauseList:
     """Create matchers for the device_ids in the event_data."""
     return DEVICE_ID_IN_EVENT.is_not(None) & sqlalchemy.cast(
         DEVICE_ID_IN_EVENT, sqlalchemy.Text()