diff --git a/homeassistant/components/logbook/queries/devices.py b/homeassistant/components/logbook/queries/devices.py index 20b56ef8dd6..5e7827b87a0 100644 --- a/homeassistant/components/logbook/queries/devices.py +++ b/homeassistant/components/logbook/queries/devices.py @@ -8,7 +8,7 @@ from sqlalchemy import Column, lambda_stmt, select, union_all from sqlalchemy.orm import Query from sqlalchemy.sql.elements import ClauseList from sqlalchemy.sql.lambdas import StatementLambdaElement -from sqlalchemy.sql.selectable import Select +from sqlalchemy.sql.selectable import CTE, CompoundSelect from homeassistant.components.recorder.models import Events, States @@ -28,7 +28,7 @@ def _select_device_id_context_ids_sub_query( end_day: dt, event_types: tuple[str, ...], json_quotable_device_ids: list[str], -) -> Select: +) -> CompoundSelect: """Generate a subquery to find context ids for multiple devices.""" return select( union_all( @@ -45,17 +45,17 @@ def _apply_devices_context_union( end_day: dt, event_types: tuple[str, ...], json_quotable_device_ids: list[str], -) -> StatementLambdaElement: +) -> CompoundSelect: """Generate a CTE to find the device context ids and a query to find linked row.""" - devices_cte = _select_device_id_context_ids_sub_query( + devices_cte: CTE = _select_device_id_context_ids_sub_query( start_day, end_day, event_types, json_quotable_device_ids, ).cte() return query.union_all( - select_events_context_only().where(Events.context_id.in_(devices_cte)), - select_states_context_only().where(States.context_id.in_(devices_cte)), + select_events_context_only().where(Events.context_id.in_(devices_cte.select())), + select_states_context_only().where(States.context_id.in_(devices_cte.select())), ) diff --git a/homeassistant/components/logbook/queries/entities.py b/homeassistant/components/logbook/queries/entities.py index 6db0931e9f3..844890c23a9 100644 --- a/homeassistant/components/logbook/queries/entities.py +++ b/homeassistant/components/logbook/queries/entities.py @@ -8,7 +8,7 @@ import sqlalchemy from sqlalchemy import Column, lambda_stmt, select, union_all from sqlalchemy.orm import Query from sqlalchemy.sql.lambdas import StatementLambdaElement -from sqlalchemy.sql.selectable import Select +from sqlalchemy.sql.selectable import CTE, CompoundSelect from homeassistant.components.recorder.models import ( ENTITY_ID_LAST_UPDATED_INDEX, @@ -37,7 +37,7 @@ def _select_entities_context_ids_sub_query( event_types: tuple[str, ...], entity_ids: list[str], json_quotable_entity_ids: list[str], -) -> Select: +) -> CompoundSelect: """Generate a subquery to find context ids for multiple entities.""" return select( union_all( @@ -58,9 +58,9 @@ def _apply_entities_context_union( event_types: tuple[str, ...], entity_ids: list[str], json_quotable_entity_ids: list[str], -) -> StatementLambdaElement: +) -> CompoundSelect: """Generate a CTE to find the entity and device context ids and a query to find linked row.""" - entities_cte = _select_entities_context_ids_sub_query( + entities_cte: CTE = _select_entities_context_ids_sub_query( start_day, end_day, event_types, @@ -69,10 +69,12 @@ def _apply_entities_context_union( ).cte() return query.union_all( states_query_for_entity_ids(start_day, end_day, entity_ids), - select_events_context_only().where(Events.context_id.in_(entities_cte)), + select_events_context_only().where( + Events.context_id.in_(entities_cte.select()) + ), select_states_context_only() .where(States.entity_id.not_in(entity_ids)) - .where(States.context_id.in_(entities_cte)), + .where(States.context_id.in_(entities_cte.select())), ) @@ -84,7 +86,6 @@ def entities_stmt( json_quotable_entity_ids: list[str], ) -> StatementLambdaElement: """Generate a logbook query for multiple entities.""" - assert json_quotable_entity_ids is not None return lambda_stmt( lambda: _apply_entities_context_union( select_events_without_states(start_day, end_day, event_types).where( diff --git a/homeassistant/components/logbook/queries/entities_and_devices.py b/homeassistant/components/logbook/queries/entities_and_devices.py index d8a23635ad7..d1c86ddbec5 100644 --- a/homeassistant/components/logbook/queries/entities_and_devices.py +++ b/homeassistant/components/logbook/queries/entities_and_devices.py @@ -8,7 +8,7 @@ import sqlalchemy from sqlalchemy import lambda_stmt, select, union_all from sqlalchemy.orm import Query from sqlalchemy.sql.lambdas import StatementLambdaElement -from sqlalchemy.sql.selectable import Select +from sqlalchemy.sql.selectable import CTE, CompoundSelect from homeassistant.components.recorder.models import Events, States @@ -33,7 +33,7 @@ def _select_entities_device_id_context_ids_sub_query( entity_ids: list[str], json_quotable_entity_ids: list[str], json_quotable_device_ids: list[str], -) -> Select: +) -> CompoundSelect: """Generate a subquery to find context ids for multiple entities and multiple devices.""" return select( union_all( @@ -57,8 +57,8 @@ def _apply_entities_devices_context_union( entity_ids: list[str], json_quotable_entity_ids: list[str], json_quotable_device_ids: list[str], -) -> StatementLambdaElement: - devices_entities_cte = _select_entities_device_id_context_ids_sub_query( +) -> CompoundSelect: + devices_entities_cte: CTE = _select_entities_device_id_context_ids_sub_query( start_day, end_day, event_types, @@ -68,10 +68,12 @@ def _apply_entities_devices_context_union( ).cte() return query.union_all( states_query_for_entity_ids(start_day, end_day, entity_ids), - select_events_context_only().where(Events.context_id.in_(devices_entities_cte)), + select_events_context_only().where( + Events.context_id.in_(devices_entities_cte.select()) + ), select_states_context_only() .where(States.entity_id.not_in(entity_ids)) - .where(States.context_id.in_(devices_entities_cte)), + .where(States.context_id.in_(devices_entities_cte.select())), )