switch to clickhouse

app.py

@@ -9,12 +9,16 @@ import os
 import torch
 from chronos import ChronosPipeline
 import numpy as np
+import requests
+import tempfile
 
 try:
     from google.colab import userdata
     PG_PASSWORD = userdata.get('FASHION_PG_PASS')
+    CH_PASSWORD = userdata.get('FASHION_CH_PASS')
 except:
     PG_PASSWORD = os.environ['FASHION_PG_PASS']
+    CH_PASSWORD = os.environ['FASHION_CH_PASS']
 
 logging.getLogger("prophet").setLevel(logging.WARNING)
 logging.getLogger("cmdstanpy").setLevel(logging.WARNING)
@@ -65,16 +69,33 @@ def read_and_process_file(file):
     return data, period_type, period_col
 
 def get_data_from_db(query):
-    conn = psycopg2.connect(
-        dbname="kroyscappingdb",
-        user="read_only",
-        password=PG_PASSWORD,
-        host="rc1d-vbh2dw5ha0gpsazk.mdb.yandexcloud.net",
-        port="6432",
-        sslmode="require"
-    )
-    data = pd.read_sql_query(query, conn)
-    conn.close()
+    # conn = psycopg2.connect(
+    #     dbname="kroyscappingdb",
+    #     user="read_only",
+    #     password=PG_PASSWORD,
+    #     host="rc1d-vbh2dw5ha0gpsazk.mdb.yandexcloud.net",
+    #     port="6432",
+    #     sslmode="require"
+    # )
+    cert_data = requests.get('https://storage.yandexcloud.net/cloud-certs/RootCA.pem').text
+
+    with tempfile.NamedTemporaryFile(delete=False) as temp_cert_file:
+        temp_cert_file.write(cert_data.encode())
+        cert_file_path = temp_cert_file.name
+
+    client = Client(host='rc1d-a93v7vf0pjfr6e2o.mdb.yandexcloud.net',
+                    port = 9440,
+                    user='user1',
+                    password=CH_PASSWORD,
+                    database='db1',
+                    secure=True,
+                    ca_certs=cert_file_path)
+
+    # data = pd.read_sql_query(query, conn)
+    result, columns = client.execute(query, with_column_types=True)
+    column_names = [col[0] for col in columns]
+    data = pd.DataFrame(result, columns=column_names)
+    # conn.close()
     return data
 
 def forecast_time_series(file, product_name, wb, ozon, model_choice):
@@ -86,11 +107,23 @@ def forecast_time_series(file, product_name, wb, ozon, model_choice):
         if ozon:
             marketplaces.append('ozon')
         mp_filter = "', '".join(marketplaces)
+        # query = f"""
+        # select
+        #     to_char(dm.end_date, 'yyyy-mm-dd') as ds,
+        #     1.0*sum(turnover) / (max(sum(turnover)) over ()) as y
+        # from v_datamart dm
+        # where {product_name}
+        #   and mp in ('{mp_filter}')
+        # group by ds
+        # order by ds
+        # """
         query = f"""
         select
-            to_char(dm.end_date, 'yyyy-mm-dd') as ds,
+            cast(start_date as date) as ds,
             1.0*sum(turnover) / (max(sum(turnover)) over ()) as y
-        from v_datamart dm
+        from datamart_all_1
+            join week_data 
+            using (id_week)
         where {product_name}
           and mp in ('{mp_filter}')
         group by ds
@@ -98,6 +131,7 @@ def forecast_time_series(file, product_name, wb, ozon, model_choice):
         """
         print(query)
         data = get_data_from_db(query)
+        print(data)
         period_type = "Week"
         period_col = "ds"
 

requirements.txt

@@ -1,7 +1,7 @@
 numpy == 1.26.4
 gradio
-pandas
+pandas == 2.1.4
 plotly
 prophet
-psycopg2
+clickhouse-driver
 git+https://github.com/amazon-science/chronos-forecasting.git