first commit

clickhouse_task/create_table.py

@@ -0,0 +1,54 @@
+import polars as pl
+from sqlalchemy import text
+
+from log import *
+
+
+def create_clickhouse_table(
+    df: pl.DataFrame,
+    table_name: str,
+    clickhouse_engine
+):
+    type_mapping = {
+        pl.Int8: "Nullable(Int8)",
+        pl.Int16: "Nullable(Int16)",
+        pl.Int32: "Nullable(Int32)",
+        pl.Int64: "Nullable(Int64)",
+        pl.UInt8: "Nullable(UInt8)",
+        pl.UInt16: "Nullable(UInt16)",
+        pl.UInt32: "Nullable(UInt32)",
+        pl.UInt64: "Nullable(UInt64)",
+        pl.Float32: "Nullable(Float32)",
+        pl.Float64: "Nullable(Float64)",
+        pl.Boolean: "Nullable(Bool)",
+        pl.String: "Nullable(String)",
+        pl.Date: "Nullable(Date)",
+        pl.Datetime: "Nullable(DateTime)",
+    }
+
+    columns = []
+
+    for col_name, dtype in df.schema.items():
+
+        clickhouse_type = type_mapping.get(
+            dtype,
+            "Nullable(String)"
+        )
+
+        columns.append(
+            f"`{col_name}` {clickhouse_type}"
+        )
+
+    create_sql = f"""
+    CREATE TABLE IF NOT EXISTS {table_name}
+    (
+        {', '.join(columns)}
+    )
+    ENGINE = MergeTree()
+    ORDER BY tuple()
+    """
+
+    with clickhouse_engine.begin() as conn:
+        conn.execute(text(create_sql))
+
+    log.info(f"Table ready: {table_name}")

clickhouse_task/delete_task.py

@@ -0,0 +1,31 @@
+from db_con.connection import *
+
+from log import *
+
+
+def truncate_table(client, table_name: str) -> None:
+    """
+    Truncate a ClickHouse table.
+    """
+    query = f"TRUNCATE TABLE {table_name}"
+
+    print(f"Truncating table: {table_name}")
+    client.command(query)
+
+    log.info(f"Table {table_name} truncated successfully.")
+
+
+
+def delete_rows(client, table_name: str, condition: str) -> None:
+    """
+    Delete rows from a ClickHouse table based on a condition.
+    """
+    query = f"""
+    ALTER TABLE {table_name}
+    DELETE WHERE {condition}
+    """
+
+    print(f"Deleting rows from {table_name} where {condition}")
+    client.command(query)
+
+    log.info("Delete command submitted successfully.")