Update ShiftManagement.py

datamodels/ShiftManagement.py

@@ -1,4 +1,7 @@
 import datetime
+import asyncio
+import logging
+from typing import Optional
 
 import aiohttp
 from bson import ObjectId
@@ -14,6 +17,7 @@ class BreakItem(BaseDataClass):
     start_epoch: int
     end_epoch: int
 
+
 class ShiftItem:
     id: str
     username: str
@@ -32,11 +36,13 @@ def __init__(self, **kwargs):
         for key, value in kwargs.items():
             setattr(self, key, value)
 
+
 class ShiftManagement:
     def __init__(self, connection, current_shifts):
         self.shifts = Document(connection, current_shifts)
+        self.logger = logging.getLogger(__name__)
 
-    async def fetch_shift(self, object_id: ObjectId) -> ShiftItem | None:
+    async def fetch_shift(self, object_id: ObjectId) -> Optional[ShiftItem]:
         shift = await self.shifts.find_by_id(object_id)
         if not shift:
             return None
@@ -57,34 +63,22 @@ async def fetch_shift(self, object_id: ObjectId) -> ShiftItem | None:
 
     async def add_shift_by_user(
         self, member: discord.Member, shift_type: str, breaks: list, guild: int, timestamp: int = 0
-    ):
+    ) -> ObjectId:
         """
-        Adds a shift for the specified user to the database, with the provided
-        extras data.
-
-        The shift is recorded as a document in the 'shifts' collection, and the
-        user's ID is used as the document ID. If a document with that ID already
-        exists in the collection, the new shift data is added to the existing
-        'data' array.
-
-        {
-          "Username": "1FriendlyDoge",
-          "Nickname": "NoobyNoob",
-          "UserID": 123456789012345678,
-          "Type": "Ingame Shift",
-          "StartEpoch": 706969420,
-          "Breaks": [
-            {
-              "StartEpoch": 706969430,
-              "EndEpoch": 706969550
-            }
-          ],
-          "Moderations": [
-                ObjectId("123456789012345678901234")
-          ],
-          "EndEpoch": 706969420,
-          "Guild": 12345678910111213
-        }
+        Adds a shift for the specified user to the database and syncs with external APIs.
+        
+        Args:
+            member: Discord member starting the shift
+            shift_type: Type of shift being started
+            breaks: List of break periods
+            guild: Guild ID where the shift is being started
+            timestamp: Optional custom start timestamp
+        
+        Returns:
+            ObjectId of the created shift document
+        
+        Raises:
+            aiohttp.ClientError: If API sync fails
         """
         data = {
             "_id": ObjectId(),
@@ -100,28 +94,47 @@ async def add_shift_by_user(
             "RemovedTime": 0,
             "EndEpoch": 0,
         }
+
         await self.shifts.db.insert_one(data)
-
+
+        url_var = config("BASE_API_URL")
+        panel_url_var = config("PANEL_API_URL")
+
+        async def sync_with_apis():
+            async with aiohttp.ClientSession() as session:
+                tasks = []
+
+                if url_var not in ["", None]:
+                    tasks.append(
+                        session.get(
+                            f"{url_var}/Internal/SyncStartShift/{data['_id']}", 
+                            headers={"Authorization": config('INTERNAL_API_AUTH')},
+                            raise_for_status=True
+                        )
+                    )
+
+                if panel_url_var not in ["", None]:
+                    tasks.append(
+                        session.post(
+                            f"{panel_url_var}/{guild}/SyncStartShift?ID={data['_id']}", 
+                            headers={"X-Static-Token": config('PANEL_STATIC_AUTH')},
+                            raise_for_status=True
+                        )
+                    )
+
+                if tasks:
+                    responses = await asyncio.gather(*tasks, return_exceptions=True)
+                    for response in responses:
+                        if isinstance(response, Exception):
+                            self.logger.error(f"API sync failed: {str(response)}")
+
         try:
-            url_var = config("BASE_API_URL")
-            panel_url_var = config("PANEL_API_URL")
-            if url_var not in ["", None]:
-                async with aiohttp.ClientSession() as session:
-                    async with session.get(
-                            f"{url_var}/Internal/SyncStartShift/{data['_id']}", headers={
-                                "Authorization": config('INTERNAL_API_AUTH')
-                            }):
-                        pass
-            if panel_url_var not in ["", None]:
-                async with aiohttp.ClientSession() as session:
-                    async with session.post(
-                            f"{panel_url_var}/{guild_id}/SyncStartShift?ID={data['_id']}", headers={
-                                "X-Static-Token": config('PANEL_STATIC_AUTH')
-                            }):
-                        pass
-        except:
-            pass
-
+            await sync_with_apis()
+        except aiohttp.ClientError as e:
+            self.logger.error(f"Failed to sync shift start with APIs: {str(e)}")
+        except Exception as e:
+            self.logger.error(f"Unexpected error during API sync: {str(e)}")
+
         return data["_id"]
 
     async def add_time_to_shift(self, identifier: str, seconds: int):
@@ -144,51 +157,87 @@ async def remove_time_from_shift(self, identifier: str, seconds: int):
 
     async def end_shift(self, identifier: str, guild_id: int | None = None, timestamp: int | None = None):
         """
-        Ends the specified user's shift.
+        Ends the specified user's shift and syncs with external APIs.
+        
+        Args:
+            identifier: Shift document ID
+            guild_id: Optional guild ID override
+            timestamp: Optional custom end timestamp
+        
+        Returns:
+            Updated shift document
+            
+        Raises:
+            ValueError: If shift not found or guild mismatch
         """
-
         document = await self.shifts.db.find_one({"_id": ObjectId(identifier)})
         if not document:
             raise ValueError("Shift not found.")
-
-        # Extract guild_id from the document if not provided
+
         guild_id = guild_id if guild_id else document["Guild"]
-
+        
         if document["Guild"] != guild_id:
             raise ValueError("Shift not found.")
-
-        document["EndEpoch"] = datetime.datetime.now().timestamp()
-
+
+        current_time = datetime.datetime.now().timestamp() if timestamp in [None, 0] else timestamp
+        document["EndEpoch"] = current_time
+
+        # Close any open breaks
         for breaks in document["Breaks"]:
             if breaks["EndEpoch"] == 0:
-                breaks["EndEpoch"] = int(datetime.datetime.now().timestamp()) if timestamp in [None, 0] else timestamp
-
+                breaks["EndEpoch"] = int(current_time)
+
+        url_var = config("BASE_API_URL")
+        panel_url_var = config("PANEL_API_URL")
+
+        async def sync_end_with_apis():
+            async with aiohttp.ClientSession() as session:
+                tasks = []
+
+                if url_var not in ["", None]:
+                    tasks.append(
+                        session.get(
+                            f"{url_var}/Internal/SyncEndShift/{document['UserID']}/{guild_id}",
+                            headers={"Authorization": config('INTERNAL_API_AUTH')},
+                            raise_for_status=True
+                        )
+                    )
+
+                if panel_url_var not in ["", None]:
+                    tasks.append(
+                        session.delete(
+                            f"{panel_url_var}/{guild_id}/SyncEndShift?ID={document['_id']}",
+                            headers={"X-Static-Token": config('PANEL_STATIC_AUTH')},
+                            raise_for_status=True
+                        )
+                    )
+
+                if tasks:
+                    responses = await asyncio.gather(*tasks, return_exceptions=True)
+                    for response in responses:
+                        if isinstance(response, Exception):
+                            self.logger.error(f"End shift API sync failed: {str(response)}")
+
         try:
-            url_var = config("BASE_API_URL")
-            panel_url_var = config("PANEL_API_URL")
-            if url_var not in ["", None]:
-                async with aiohttp.ClientSession() as session:
-                    async with session.get(
-                            f"{url_var}/Internal/SyncEndShift/{document['UserID']}/{guild_id}", headers={
-                                "Authorization": config('INTERNAL_API_AUTH')
-                            }):
-                        pass
-            if panel_url_var not in ["", None]:
-                async with aiohttp.ClientSession() as session:
-                    async with session.delete(
-                            f"{panel_url_var}/{guild_id}/SyncEndShift?ID={document['_id']}", headers={
-                                "X-Static-Token": config('PANEL_STATIC_AUTH')
-                            }):
-                        pass
-        except:
-            pass
-
+            await sync_end_with_apis()
+        except aiohttp.ClientError as e:
+            self.logger.error(f"Failed to sync shift end with APIs: {str(e)}")
+        except Exception as e:
+            self.logger.error(f"Unexpected error during end shift API sync: {str(e)}")
+
         await self.shifts.update_by_id(document)
         return document
 
     async def get_current_shift(self, member: discord.Member, guild_id: int):
         """
         Gets the current shift for the specified user.
+        
+        Args:
+            member: Discord member to check
+            guild_id: Guild ID to check
+            
+        Returns:
+            Current shift document or None if no active shift
         """
         return await self.shifts.db.find_one(
             {"UserID": member.id, "EndEpoch": 0, "Guild": guild_id}