Update main.py

main.py

@@ -250,10 +250,15 @@ async def generate_product_count_prediction(b_id: int):
 
 @app.post("/get_product_count_prediction_from_DB")
 async def get_product_count_prediction_from_DB(b_id: int,page_number :int,page_size :int):
+        try:
+            data = get_data(b_id,page_number,page_size,False)
+        except Exception as e:
+            data = str(e)
+            print(str(e))
         response_content = {
             "status": "done",
             "message": "data from DB",
-            "data": get_data(b_id,page_number,page_size,False),
+            "data": data,
             "status_code":200
         }
         return response_content
@@ -276,8 +281,12 @@ async def get_product_count_prediction_from_DB(
                                                search_sku: str = Query(None, description="The SKU of the product to search"),
                                                product_name: str = Query(None, description="The name of the product to search")
                                             ):
+        try:
+            related_products = find_related_products(sub_sku=search_sku, product_name=product_name,forecast_data = get_data(b_id,1,2,True))
+        except Exception as e:
+            related_products = str(e)
+            print(str(e))
     
-        related_products = find_related_products(sub_sku=search_sku, product_name=product_name,forecast_data = get_data(b_id,1,2,True))
         response_content = {
             "status": "done",
             "message": "data from DB",