Improve performance for history component

homeassistant/components/history.py

@@ -7,6 +7,7 @@ Provide pre-made queries on top of the recorder component.
 import re
 from datetime import datetime, timedelta
 from itertools import groupby
+from collections import defaultdict
 
 import homeassistant.components.recorder as recorder
 
@@ -30,9 +31,7 @@ def last_5_states(entity_id):
 def state_changes_during_period(start_time, end_time=None, entity_id=None):
     """
     Return states changes during period start_time - end_time.
-    Currently does _not_ include how the states where at exactly start_time.
     """
-
     where = "last_changed=last_updated AND last_changed > ? "
     data = [start_time]
 
@@ -49,24 +48,52 @@ def state_changes_during_period(start_time, end_time=None, entity_id=None):
 
     states = recorder.query_states(query, data)
 
-    result = []
+    result = defaultdict(list)
 
+    # Get the states at the start time
+    for state in get_states(start_time):
+        state.last_changed = start_time
+        result[state.entity_id].append(state)
+
+    # Append all changes to it
     for entity_id, group in groupby(states, lambda state: state.entity_id):
-        # Query the state of the entity ID before `start_time` so the returned
+        result[entity_id].extend(group)
-        # set will cover everything between `start_time` and `end_time`.
-        old_state = list(recorder.query_states(
-            "SELECT * FROM states WHERE entity_id = ? AND last_changed <= ? "
-            "AND last_changed=last_updated ORDER BY last_changed DESC "
-            "LIMIT 0, 1", (entity_id, start_time)))
-
-        if old_state:
-            old_state[0].last_changed = start_time
-
-        result.append(old_state + list(group))
 
     return result
 
 
+def get_states(point_in_time, entity_ids=None, run=None):
+    """ Returns the states at a specific point in time. """
+    if run is None:
+        run = recorder.run_information(point_in_time)
+
+    where = run.where_after_start_run + "AND created < ? "
+    where_data = [point_in_time]
+
+    if entity_ids is not None:
+        where += "AND entity_id IN ({}) ".format(
+            ",".join(['?'] * len(entity_ids)))
+        where_data.extend(entity_ids)
+
+    query = """
+        SELECT * FROM states
+        INNER JOIN (
+            SELECT max(state_id) AS max_state_id
+            FROM states WHERE {}
+            GROUP BY entity_id)
+        WHERE state_id = max_state_id
+    """.format(where)
+
+    return recorder.query_states(query, where_data)
+
+
+def get_state(point_in_time, entity_id, run=None):
+    """ Return a state at a specific point in time. """
+    states = get_states(point_in_time, (entity_id,), run)
+
+    return states[0] if states else None
+
+
 def setup(hass, config):
     """ Setup history hooks. """
     hass.http.register_path(
@@ -87,7 +114,7 @@ def _api_last_5_states(handler, path_match, data):
     """ Return the last 5 states for an entity id as JSON. """
     entity_id = path_match.group('entity_id')
 
-    handler.write_json(list(last_5_states(entity_id)))
+    handler.write_json(last_5_states(entity_id))
 
 
 def _api_history_period(handler, path_match, data):
@@ -98,4 +125,4 @@ def _api_history_period(handler, path_match, data):
     entity_id = data.get('filter_entity_id')
 
     handler.write_json(
-        state_changes_during_period(start_time, entity_id=entity_id))
+        state_changes_during_period(start_time, entity_id=entity_id).values())