14-06-26 1st comit

test2.py

@@ -0,0 +1,275 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#   "polars>=0.20.0",
+#   "pyarrow>=18.0.0",
+#   "sqlalchemy>=2.0.0",
+#   "pyodbc>=5.0.0",
+#   "clickhouse-connect>=0.7.0",
+#   "clickhouse-sqlalchemy>=0.3.2",
+#   "pyyaml>=6.0.3",
+#   "python-dotenv>=1.0.0",
+# ]
+# ///
+
+
+from __future__ import annotations
+
+import sys
+from datetime import date, datetime, timedelta
+
+import polars as pl
+import yaml
+
+
+from log import log
+
+from clickhouse_task.create_table import create_clickhouse_table
+from clickhouse_task.delete_task import (
+    delete_existing_data,
+    truncate_table,
+)
+
+from clickhouse_task.load_table import load_to_clickhouse
+
+from db_con.connection import (
+    build_sql_server_engine,
+    build_clickhouse_engine,
+    get_clickhouse_client,
+)
+
+from mids import (
+    MID_TABLE_COV,
+    MID_TABLE_COV1,
+)
+
+from masters.dimensions import *
+from masters.bridge import *
+from kpi.facts import *
+
+
+# ==========================================================
+# Helpers
+# ==========================================================
+
+def table_exists(
+    client,
+    table_name: str,
+) -> bool:
+
+    return bool(
+        client.command(
+            f"EXISTS TABLE {table_name}"
+        )
+    )
+
+
+def get_dataframe(
+    fn_name: str,
+    fetch_by: str,
+    sql_engine,
+    mids,
+    run_date,
+):
+
+    fn = globals()[fn_name]
+
+    if fetch_by == "mids":
+        return fn(sql_engine, mids)
+
+    if fetch_by == "run_date":
+        return fn(sql_engine, run_date)
+
+    return fn(sql_engine)
+
+
+# ==========================================================
+# Main
+# ==========================================================
+
+def main():
+
+    log.info("=" * 80)
+    log.info("Hello from data-move Python data pipeline!")
+
+    # ------------------------------------------------------
+    # Run Date
+    # ------------------------------------------------------
+
+    if len(sys.argv) > 1:
+        run_date = datetime.strptime(
+            sys.argv[1],
+            "%Y-%m-%d",
+        ).date()
+    else:
+        run_date = date.today() - timedelta(days=1)
+
+    log.info(
+        "Pipeline Run Date: %s",
+        run_date,
+    )
+
+    # ------------------------------------------------------
+    # Connections
+    # ------------------------------------------------------
+
+    log.info(
+        "Connecting to databases..."
+    )
+
+    sql_engine = build_sql_server_engine()
+    clickhouse_engine = build_clickhouse_engine()
+    client = get_clickhouse_client()
+
+    log.info(
+        "Database connections established"
+    )
+
+    # ------------------------------------------------------
+    # Delete Keys
+    # ------------------------------------------------------
+
+    mids = MID_TABLE_COV(
+        sql_engine,
+        run_date,
+    )
+
+    emp_visit_df = MID_TABLE_COV1(
+        sql_engine,
+        run_date,
+    )
+
+    # ------------------------------------------------------
+    # Config
+    # ------------------------------------------------------
+
+    with open(
+        "tables.yml",
+        "r",
+    ) as file:
+
+        config = yaml.safe_load(file)
+
+    # ------------------------------------------------------
+    # Process Tables
+    # ------------------------------------------------------
+
+    for table in config["tables"]:
+
+        table_name = table["name"]
+        operation = table["operation"]
+        fetch_by = table["fetch_by"]
+
+        log.info("=" * 80)
+        log.info(
+            "Processing Table: %s",
+            table_name,
+        )
+
+        try:
+
+            # ------------------------------------------
+            # Fetch Data
+            # ------------------------------------------
+
+            fn_name = f"fetch_{table_name}"
+
+            df = get_dataframe(
+                fn_name=fn_name,
+                fetch_by=fetch_by,
+                sql_engine=sql_engine,
+                mids=mids,
+                run_date=run_date,
+            )
+
+            if df.is_empty():
+
+                log.warning(
+                    "%s returned no rows",
+                    table_name,
+                )
+
+                continue
+
+            log.info(
+                "Fetched %s rows",
+                len(df),
+            )
+
+            # ------------------------------------------
+            # Create Table If Missing
+            # ------------------------------------------
+
+            exists = table_exists(
+                client,
+                table_name,
+            )
+
+            if not exists:
+
+                log.info(
+                    "Creating table %s",
+                    table_name,
+                )
+
+                create_clickhouse_table(
+                    df=df,
+                    table_name=table_name,
+                    clickhouse_engine=clickhouse_engine,
+                )
+
+            # ------------------------------------------
+            # Existing Table Logic
+            # ------------------------------------------
+
+            else:
+
+                if operation == "DELETE+INSERT":
+
+                    truncate_table(
+                        client,
+                        table_name,
+                    )
+
+                else:
+
+                    delete_existing_data(
+                        client=client,
+                        table_name=table_name,
+                        run_date=run_date,
+                        mids=mids,
+                        emp_visit_df=emp_visit_df,
+                    )
+
+            # ------------------------------------------
+            # Load Data
+            # ------------------------------------------
+
+            load_to_clickhouse(
+                client=client,
+                table_name=table_name,
+                df=df,
+            )
+
+            log.info(
+                "%s loaded successfully (%s rows)",
+                table_name,
+                len(df),
+            )
+
+        except Exception:
+
+            log.exception(
+                "Failed processing table %s",
+                table_name,
+            )
+
+            raise
+
+    log.info("=" * 80)
+    log.info("Pipeline Completed Successfully")
+    log.info("=" * 80)
+
+
+if __name__ == "__main__":
+    main()