fix potential cursor issue

Author: zzstoatzz

server/data_stream.py

@@ -1,7 +1,14 @@
 import logging
 from collections import defaultdict
 
-from atproto import AtUri, CAR, firehose_models, FirehoseSubscribeReposClient, models, parse_subscribe_repos_message
+from atproto import (
+    AtUri,
+    CAR,
+    firehose_models,
+    FirehoseSubscribeReposClient,
+    models,
+    parse_subscribe_repos_message,
+)
 from atproto.exceptions import FirehoseError
 
 from server.database import SubscriptionState
@@ -15,21 +22,21 @@
 
 
 def _get_ops_by_type(commit: models.ComAtprotoSyncSubscribeRepos.Commit) -> defaultdict:
-    operation_by_type = defaultdict(lambda: {'created': [], 'deleted': []})
+    operation_by_type = defaultdict(lambda: {"created": [], "deleted": []})
 
     car = CAR.from_bytes(commit.blocks)
     for op in commit.ops:
-        if op.action == 'update':
+        if op.action == "update":
             # we are not interested in updates
             continue
 
-        uri = AtUri.from_str(f'at://{commit.repo}/{op.path}')
+        uri = AtUri.from_str(f"at://{commit.repo}/{op.path}")
 
-        if op.action == 'create':
+        if op.action == "create":
             if not op.cid:
                 continue
 
-            create_info = {'uri': str(uri), 'cid': str(op.cid), 'author': commit.repo}
+            create_info = {"uri": str(uri), "cid": str(op.cid), "author": commit.repo}
 
             record_raw_data = car.blocks.get(op.cid)
             if not record_raw_data:
@@ -40,12 +47,16 @@ def _get_ops_by_type(commit: models.ComAtprotoSyncSubscribeRepos.Commit) -> defa
                 continue
 
             for record_type, record_nsid in _INTERESTED_RECORDS.items():
-                if uri.collection == record_nsid and models.is_record_type(record, record_type):
-                    operation_by_type[record_nsid]['created'].append({'record': record, **create_info})
+                if uri.collection == record_nsid and models.is_record_type(
+                    record, record_type
+                ):
+                    operation_by_type[record_nsid]["created"].append(
+                        {"record": record, **create_info}
+                    )
                     break
 
-        if op.action == 'delete':
-            operation_by_type[uri.collection]['deleted'].append({'uri': str(uri)})
+        if op.action == "delete":
+            operation_by_type[uri.collection]["deleted"].append({"uri": str(uri)})
 
     return operation_by_type
 
@@ -57,7 +68,7 @@ def run(name, operations_callback, stream_stop_event=None):
         except FirehoseError as e:
             if logger.level == logging.DEBUG:
                 raise e
-            logger.error(f'Firehose error: {e}. Reconnecting to the firehose.')
+            logger.error(f"Firehose error: {e}. Reconnecting to the firehose.")
 
 
 def _run(name, operations_callback, stream_stop_event=None):
@@ -69,9 +80,6 @@ def _run(name, operations_callback, stream_stop_event=None):
 
     client = FirehoseSubscribeReposClient(params)
 
-    if not state:
-        SubscriptionState.create(service=name, cursor=0)
-
     def on_message_handler(message: firehose_models.MessageFrame) -> None:
         # stop on next message if requested
         if stream_stop_event and stream_stop_event.is_set():
@@ -84,9 +92,15 @@ def on_message_handler(message: firehose_models.MessageFrame) -> None:
 
         # update stored state every ~1k events
         if commit.seq % 1000 == 0:  # lower value could lead to performance issues
-            logger.debug(f'Updated cursor for {name} to {commit.seq}')
-            client.update_params(models.ComAtprotoSyncSubscribeRepos.Params(cursor=commit.seq))
-            SubscriptionState.update(cursor=commit.seq).where(SubscriptionState.service == name).execute()
+            logger.debug(f"Updated cursor for {name} to {commit.seq}")
+            client.update_params(
+                models.ComAtprotoSyncSubscribeRepos.Params(cursor=commit.seq)
+            )
+            SubscriptionState.insert(service=name, cursor=commit.seq).on_conflict(
+                conflict_target=(SubscriptionState.service,),
+                action="UPDATE",
+                update={SubscriptionState.cursor: commit.seq},
+            ).execute()
 
         if not commit.blocks:
             return