🤖 Refactor Inventory Management and Add Atomic Update Functionality

flask/src/db.py

@@ -127,25 +127,20 @@ def get_products_join():
 
 def get_inventory(cart):
     print("> get_inventory")
-
     quantities = cart['quantities']
-
     print("> quantities", quantities)
-
-    productIds = []
-    for productId in quantities:
-        productIds.append(productId)
-
-    productIds = formatArray(productIds)
-    print("> productIds", productIds)
-
     try:
+        productIds = list(quantities.keys())
+
         with sentry_sdk.start_span(op="get_inventory", description="db.connect"):
             connection = db.connect()
         with sentry_sdk.start_span(op="get_inventory", description="db.query") as span:
+            query = text("SELECT id, productId, count FROM inventory WHERE productId = ANY(:product_ids)")
             inventory = connection.execute(
-                "SELECT * FROM inventory WHERE productId in %s" % (productIds)
+                query,
+                {"product_ids": productIds}
             ).fetchall()
+            inventory = [dict(row) for row in inventory]
             span.set_data("inventory",inventory)
     except BrokenPipeError as err:
         raise DatabaseConnectionError('get_inventory')
@@ -158,6 +153,26 @@ def get_inventory(cart):
 
     return inventory
 
+def update_inventory_quantity(product_id, requested_quantity, connection=None):
+    """Updates inventory count atomically after successful checkout."""
+    should_close = False
+    try:
+        if not connection:
+            connection = db.connect()
+            should_close = True
+
+        query = text("""
+            UPDATE inventory 
+            SET count = count - :quantity 
+            WHERE productId = :product_id AND count >= :quantity
+            RETURNING count""")
+
+        result = connection.execute(query, {"product_id": product_id, "quantity": requested_quantity}).fetchone()
+        return result[0] if result else None
+    finally:
+        if should_close and connection:
+            connection.close()
+
 
 
 def formatArray(ids):

flask/src/main.py

@@ -118,20 +118,47 @@ def checkout():
     except Exception as err:
         raise (err)
 
-    print("> /checkout inventory", inventory)
-    print("> validate_inventory", validate_inventory)
+    if not inventory or not quantities:
+        sentry_sdk.metrics.incr(key="checkout.failed")
+        raise Exception("No inventory found for requested products")
 
     with sentry_sdk.start_span(op="process_order", description="function"):
+        # Create lookup map for efficient inventory checks
+        inventory_map = {str(item['productId']): item for item in inventory}
         quantities = cart['quantities']
-        for cartItem in quantities:
-            for inventoryItem in inventory:
-                print("> inventoryItem.count", inventoryItem['count'])
-                if (validate_inventory and (inventoryItem.count < quantities[cartItem] or quantities[cartItem] >= inventoryItem.count)):
+
+        if validate_inventory:
+            # Validate all quantities before updating anything
+            for product_id, requested_quantity in quantities.items():
+                if str(product_id) not in inventory_map:
                     sentry_sdk.metrics.incr(key="checkout.failed")
-                    raise Exception("Not enough inventory for product")
-        if len(inventory) == 0 or len(quantities) == 0:
-            raise Exception("Not enough inventory for product")
+                    raise Exception(f"Product {product_id} not found in inventory")
+
+                inventory_item = inventory_map[str(product_id)]
+                available_quantity = inventory_item['count']
+                requested_quantity = int(requested_quantity)
 
+                if requested_quantity <= 0:
+                    sentry_sdk.metrics.incr(key="checkout.failed")
+                    raise Exception("Invalid quantity requested")
+
+                if requested_quantity > available_quantity:
+                    sentry_sdk.metrics.incr(key="checkout.failed")
+                    raise Exception(
+                        f"Not enough inventory for product {product_id}. "
+                        f"Requested: {requested_quantity}, Available: {available_quantity}"
+                    )
+
+            # Update inventory atomically
+            try:
+                connection = db.connect()
+                for product_id, requested_quantity in quantities.items():
+                    new_count = update_inventory_quantity(product_id, int(requested_quantity), connection)
+                    if new_count is None:
+                        raise Exception(f"Failed to update inventory for product {product_id}")
+            except Exception as err:
+                sentry_sdk.metrics.incr(key="checkout.failed")
+                raise Exception("Failed to update inventory: " + str(err))
     response = make_response("success")
     return response