Update main.py

main.py

@@ -17,6 +17,136 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+def insert_data(bid,forecast_data):#mysql-connector-python
+    import mysql.connector
+    import json
+    
+    # Define connection parameters
+    host = "68.183.225.237"
+    user = "sm_ml"
+    password = "Fz6/I733"
+    database = "sm_qa_1"
+    
+    # Establish connection
+    connection = mysql.connector.connect(
+        host=host,
+        user=user,
+        password=password,
+        database=database
+    )
+    # Create a cursor object
+    cursor = connection.cursor()
+    # Convert forecast_data to JSON string
+    forecast_data_json = json.dumps(forecast_data)
+    
+    # SQL command to insert data
+    insert_query = """
+    INSERT INTO sm_product_count_forecast (bid, forecast_data)
+    VALUES (%s, %s)
+    """
+    
+    # Execute the SQL command with data
+    cursor.execute(insert_query, (bid, forecast_data_json))
+    
+    # Commit the transaction
+    connection.commit()
+    
+    print("Data inserted successfully")
+    
+    # Close the cursor and connection
+    cursor.close()
+    connection.close()
+
+
+def delete_json(b_id):
+    import mysql.connector
+    
+    # Define connection parameters
+    host = "68.183.225.237"
+    user = "sm_ml"
+    password = "Fz6/I733"
+    database = "sm_qa_1"
+    
+    # Establish connection
+    connection = mysql.connector.connect(
+        host=host,
+        user=user,
+        password=password,
+        database=database
+    )
+    # Create a cursor object
+    cursor = connection.cursor()
+    # SQL command to delete a specific record
+    delete_query = """
+    DELETE FROM sm_product_count_forecast
+    WHERE bid = %s
+    """
+    
+    # Execute the SQL command with the specified BID
+    cursor.execute(delete_query, (b_id,))
+    # Commit the transaction
+    connection.commit()
+    print(f"Record with BID {bid_to_delete} deleted successfully")
+    # Close the cursor and connection
+    cursor.close()
+    connection.close()
+
+def get_data(b_id):
+    import mysql.connector
+    import json
+    
+    # Define connection parameters
+    host = "68.183.225.237"
+    user = "sm_ml"
+    password = "Fz6/I733"
+    database = "sm_qa_1"
+    
+    # Establish connection
+    connection = mysql.connector.connect(
+        host=host,
+        user=user,
+        password=password,
+        database=database
+    )
+    
+    # Create a cursor object
+    cursor = connection.cursor()
+    
+    # SQL command to select data for a specific BID
+    select_query = """
+    SELECT bid, forecast_data, created_at
+    FROM sm_product_count_forecast
+    WHERE bid = %s
+    """
+    
+    # Execute the SQL command with the specified BID
+    cursor.execute(select_query, (b_id,))
+    
+    # Fetch the result
+    row = cursor.fetchone()
+    
+    if row:
+        bid = row[0]
+        forecast_data_json = row[1]
+        created_at = row[2]
+        
+        # Convert JSON string back to Python dictionary
+        forecast_data = json.loads(forecast_data_json)
+        result = {
+            "BID":bid,
+            "created_at":created_at,
+            "forecast_data":forecast_data
+        }
+        return result
+        
+    else:
+        return None
+    
+    # Close the cursor and connection
+    cursor.close()
+    connection.close()
+
+
 def forecast(monthly_sales):
     # Prepare the data for Prophet
     monthly_sales.rename(columns={'transaction_date': 'ds', 'sell_qty': 'y'}, inplace=True)
@@ -71,19 +201,19 @@ def process_product(product_name, data):
         return {
             "Product Name": product_name,
             "next_month": str(result_dict["next_month"]),
-            "predicted_count": result_dict["predicted_count"],
-            "full_trend" : str(full_trend)
+            "predicted_count": result_dict["predicted_count"]
+            #"full_trend" : str(full_trend)
         }
     except Exception as e:
         return {
             "Product Name": product_name,
             "next_month": str(e),
-            "predicted_count": "not predicted",
-            "full_trend" : str(full_trend)
+            "predicted_count": "not predicted"
+            #"full_trend" : str(full_trend)
         }
 
-@app.post("/get_product_count_prediction")
-async def get_product_count_prediction(b_id: int):
+@app.post("/generate_product_count_prediction")
+async def generate_product_count_prediction(b_id: int):
     try:
         # main
         data, message = dc.get_data(b_id=b_id, product_name="sample")
@@ -98,20 +228,33 @@ async def get_product_count_prediction(b_id: int):
                 for future in as_completed(futures):
                     results.append(future.result())
 
-            response_content = {
-                "status": "success",
-                "message": "Prediction successful",
-                "data": results
-            }
-            return JSONResponse(content=response_content, status_code=200)
+
+            delete_json(b_id)
+            insert_data(bid,results)
+            return {"status": "success",
+                    "b_id":b_id,
+                    "message": "Prediction successful and saved to DB",
+                    "status_code":200
+                   }
             
-        else:
-            raise HTTPException(status_code=400, detail=message)
+
     except Exception as e:
         print(str(e))
         response_content = {
             "status": "error",
             "message": str(e),
-            "data": None
+            "data": None,
+            "status_code":200
+        }
+        return response_content
+
+
+@app.post("/get_product_count_prediction_from_DB")
+async def get_product_count_prediction_from_DB(b_id: int):
+        response_content = {
+            "status": "done",
+            "message": "data from DB",
+            "data": get_data(b_id),
+            "status_code":200
         }
-        return JSONResponse(content=response_content, status_code=500)
+        return response_content