stats: add lexical query to history

CHANGELOG.rst

@@ -12,6 +12,7 @@ and this project adheres to `Semantic Versioning <https://semver.org/spec/v2.0.0
 Added
 -----
 - Read version from bot
+- stats: add lexical query to history
 
 Removed
 -------

README.rst

@@ -216,6 +216,9 @@ week
 history
 -------
 ``/stats history`` returns a plot of messages versus date.
+Allows limiting by a lexical query (using Postgres'
+`tsquery syntax <https://www.postgresql.org/docs/12/datatype-textsearch.html#DATATYPE-TSQUERY>`_)
+with the ``lquery`` option.
 
 .. image:: examples/history.png
    :alt: Example of history plot

telegram_stats_bot/db.py

@@ -46,6 +46,13 @@ def init_dbs(engine: Engine):
             type                    text
         );
         
+        alter table messages_utc
+                add column if not exists text_index_col tsvector
+                generated always as (to_tsvector('english', coalesce(text, ''))) stored;
+        
+        create index if not exists text_idx
+            on messages_utc using gin (text_index_col);
+        
         create index if not exists messages_utc_date_index
             on messages_utc (date);
         

telegram_stats_bot/stats.py

@@ -406,11 +406,12 @@ class StatsRunner(object):
 
         return None, bio
 
-    def get_message_history(self, user: Tuple[int, str] = None, averages: int = None, start: str = None,
+    def get_message_history(self, user: Tuple[int, str] = None, lquery: str = None, averages: int = None, start: str = None,
                             end: str = None) \
             -> Tuple[None, BytesIO]:
         """
         Make a plot of message history over time
+        :param lquery: Limit results to lexical query (&, |, !, <n>)
         :param averages: Moving average width (in days)
         :param start: Start timestamp (e.g. 2019, 2019-01, 2019-01-01, "2019-01-01 14:21")
         :param end: End timestamp (e.g. 2019, 2019-01, 2019-01-01, "2019-01-01 14:21")
@@ -422,6 +423,9 @@ class StatsRunner(object):
             if averages < 0:
                 raise HelpException("averages must be >= 0")
 
+        if lquery:
+            query_conditions.append(f"text_index_col @@ to_tsquery('{lquery}')")
+
         if start:
             sql_dict['start_dt'] = pd.to_datetime(start)
             query_conditions.append("date >= %(start_dt)s")
@@ -462,14 +466,16 @@ class StatsRunner(object):
         else:
             alpha = 1
 
-        fig = Figure()  # TODO: One day pandas will let you use constrained_layout=True here...
+        fig = Figure(constrained_layout=True)
         subplot = fig.subplots()
         df.plot(x='day', y='messages', alpha=alpha, legend=False, ax=subplot)
         if averages:
             df.plot(x='day', y='msg_rolling', legend=False, ax=subplot)
         subplot.set_ylabel("Messages")
         subplot.set_xlabel("Date")
-        if user:
+        if lquery:
+            subplot.set_title(f"History for query: {lquery}")
+        elif user:
             subplot.set_title(f"Message History for {user[1]}")
         else:
             subplot.set_title("Message History")