quasardb 3.14.2.dev8__cp310-cp310-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl

quasardb/firehose.py

@@ -0,0 +1,112 @@
+import logging
+import time
+from typing import Any, Dict, Iterator, List, Optional, Tuple
+
+import numpy as np
+
+from quasardb import Cluster
+
+FIREHOSE_TABLE = "$qdb.firehose"
+POLL_INTERVAL = 0.1
+
+logger = logging.getLogger("quasardb.firehose")
+
+
+def _init() -> Dict[str, Any]:
+    """
+    Initialize our internal state.
+    """
+    return {"last": None, "seen": set()}
+
+
+def _get_transactions_since(
+    conn: Cluster, table_name: str, last: Optional[Dict[str, Any]]
+) -> List[Dict[str, Any]]:
+    """
+    Retrieve all transactions since a certain timestamp. `last` is expected to be a dict
+    firehose row with at least a $timestamp attached.
+    """
+    if last is None:
+        q = "SELECT $timestamp, transaction_id, begin, end FROM \"{}\" WHERE table = '{}' ORDER BY $timestamp".format(
+            FIREHOSE_TABLE, table_name
+        )
+    else:
+        q = "SELECT $timestamp, transaction_id, begin, end FROM \"{}\" IN RANGE ({}, +1y) WHERE table = '{}' ORDER BY $timestamp".format(
+            FIREHOSE_TABLE, last["$timestamp"], table_name
+        )
+
+    return conn.query(q)
+
+
+def _get_transaction_data(
+    conn: Cluster, table_name: str, begin: str, end: str
+) -> List[Dict[str, Any]]:
+    """
+    Gets all data from a certain table.
+    """
+    q = 'SELECT * FROM "{}" IN RANGE ({}, {}) '.format(table_name, begin, end)
+    return conn.query(q)
+
+
+def _get_next(
+    conn: Cluster, table_name: str, state: Dict[str, Any]
+) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
+
+    # Our flow to retrieve new data is as follows:
+    # 1. Based on the state's last processed transaction, retrieve all transactions
+    #    that are logged into the firehose since then.
+    # 2. For each of the transactions, verify we haven't seen it before
+    # 3. For each of the transactions, pull in all data
+    # 4. Concatenate all this data (in order of quasardb transaction)
+
+    txs = _get_transactions_since(conn, table_name, state["last"])
+
+    xs: List[Dict[str, Any]] = []
+    for tx in txs:
+        txid = tx["transaction_id"]
+
+        if state["last"] is not None and tx["$timestamp"] > state["last"]["$timestamp"]:
+            # At this point we are guaranteed that the transaction we encounter is
+            # 'new', will not conflict with any other transaction ids. It is thus
+            # safe to reset the txid set.
+            state["seen"] = set()
+
+        if txid not in state["seen"]:
+            xs = xs + _get_transaction_data(
+                conn,
+                table_name,
+                tx["begin"],
+                # The firehose logs transaction `end` span as
+                # end inclusive, while our bulk reader and/or query
+                # language are end exclusive.
+                tx["end"] + np.timedelta64(1, "ns"),
+            )
+
+            # Because it is possible that multiple firehose changes are stored with the
+            # exact same $timestamp, we also keep track of the actually seen
+            # transaction ids.
+            state["seen"].add(txid)
+
+        state["last"] = tx
+
+    return (state, xs)
+
+
+def subscribe(conn: Cluster, table_name: str) -> Iterator[Dict[str, Any]]:
+    state = _init()
+
+    while True:
+        # Note how this is effectively a never-ending fold loop
+        # that transforms state into a new state. This state effectively
+        # functions as a checkpoint.
+        #
+        # At a later point, we could choose to provide the user
+        # direct access to this 'state' object, so that they can
+        # implement e.g. mechanisms to replay from a certain checkpoint.
+        (state, xs) = _get_next(conn, table_name, state)
+
+        for x in xs:
+            yield x
+
+        # Our poll interval
+        time.sleep(POLL_INTERVAL)