(20250117) Day-end push.

controllers_v2/core/auth_token.py

@@ -21,7 +21,7 @@
         N/A
 
 """
-
+import datetime
 #   *****************************************************************************************************************
 #  *****                                                                                                          ****
 # ***                                                   IMPORT                                                      ***
@@ -55,6 +55,7 @@ import httpx
 
 # To work with MongoDB:
 from bson import ObjectId
+from pymongo import UpdateOne
 
 
 #   *****************************************************************************************************************
@@ -571,6 +572,117 @@ class CoreAuthTokenController(CoreBaseModel):
         # Done here:
         return [CoreAuthTokenModel(**token) for token in tokens]
 
+    async def get_batches_to_sync(
+            self,
+            mongo_data_conn: AsyncMongo,
+            limit: int = 25,
+            batch_timeout_seconds: int | float = 120,
+            additional_filter: dict = None
+    ) -> List[CoreAuthTokenModel]:
+
+        """
+        Use this when some kind of account that has been integrated needs to perform data synchronization at periodic
+        intervals. One simple example is mails. The aim is to pick batches of auth-tokens
+        :param mongo_data_conn: The database connection (MongoDB) to use to perform the action.
+        :param limit: The max. no. of records to pick.
+        :param batch_timeout_seconds: The no. of seconds to not pick an element that has been marked as a member of
+            another batch.
+        :param additional_filter: Any addition filters to use.
+        :return: The retrieved record that has the token, and information about the service and client if found, else
+            None when there is no matching record.
+        """
+
+        # Prepare the filter:
+        now = date_time.get_current_utc_date_time(as_string = False)
+        batch_timeout_ts = now - datetime.timedelta(seconds = int(batch_timeout_seconds))
+        filter_json = {
+            "$and": [
+                {
+                    "$or": [
+                        {"syncAfterTs": None},  # ............ The time after which sync'ing is allowed is null.
+                        {"syncAfterTs": {"$lte": now}},  # ... The time after which sync'ing is allowed has passed.
+                    ]
+                },
+                {
+                    "$or": [
+                        {"batchId": None},  # ....................... Not currently a member of a batch.
+                        {"batchTs": {"$lt": batch_timeout_ts}}  # ... Is a member of a batch, but it has timed-out.
+                    ]
+                }
+            ]
+        }
+        if self._base_filter:
+            for k, v in self._base_filter.items(): filter_json[k] = v
+        if additional_filter:
+            for k, v in additional_filter.items(): filter_json[k] = v
+
+        # Generate a batch-id:
+        batch_id = ObjectId()
+
+        # fetch the credentials that match the conditions:
+        tokens = await mongo_data_conn.find_many(
+            collection = self.AUTH_COLLECTION,
+            filter = filter_json,
+            sort = {"lastSyncTs": 1},
+            limit = limit,
+            projection = {"_id": True}
+        )
+        token_ids = [t["_id"] for t in tokens]
+
+        # Mark 'n' no. of auth-tokens as a member of your batch:
+        await mongo_data_conn.update_many(
+            collection = self.AUTH_COLLECTION,
+            filter = {"_id": {"$in": token_ids}},
+            update = {"$set": {"batchId": batch_id, "batchTs": now}},
+            upsert = False
+        )
+
+        # Fetch the credentials to use for this batch:
+        tokens = await mongo_data_conn.find_many(
+            collection = self.AUTH_COLLECTION,
+            filter = {"batchId": batch_id},
+            limit = len(token_ids)
+        )
+
+        # Done here:
+        return [CoreAuthTokenModel(**token) for token in tokens]
+
+    async def release_token_from_batch_by_id(
+            self,
+            mongo_data_conn: AsyncMongo,
+            token_id: ObjectId | str,
+            sync_after_ts: datetime.datetime | None,
+            last_sync_ts: datetime.datetime | None = None
+    ) -> bool:
+
+        """
+        To release a particular auth-tokn from a sync'ing batch. This helps in gracefully completing the cycle.
+        :param mongo_data_conn: The database connection (MongoDB) to use to perform the action.
+        :param token_id: The ObjectId of the document in MongoDb that holds the auth-token.
+        :param sync_after_ts: The time (UTC) after which this same token can be picked again for sync'ing.
+        :param last_sync_ts: The time at which this latest sync'ing was done. Leave it to null for the current time,
+            else pass a custom timestamp.
+        :return: True if successful, else False.
+        """
+
+        # Process the inputs:
+        if isinstance(last_sync_ts, datetime.datetime): date_time.to_timezone(last_sync_ts, date_time.TIMEZONE_UTC)
+        else: last_sync_ts = date_time.get_current_utc_date_time(as_string = False)
+
+        # release the auth-token from the batch:
+        return await mongo_data_conn.update_one(
+            collection = self.AUTH_COLLECTION,
+            filter = {"_id": ObjectId(token_id)},
+            update = {
+                "$set": {
+                    "batchId": None,
+                    "syncAfterTs": sync_after_ts,
+                    "lastSyncTs": last_sync_ts
+                }
+            },
+            upsert = False
+        )
+
     # ┏┓┳┓┳┳┳┓      ┳┓  ┓
     # ┃ ┣┫┃┃┃┃  ━━  ┃┃┏┓┃┏┓╋┏┓
     # ┗┛┛┗┗┛┻┛      ┻┛┗ ┗┗ ┗┗