Files
airflow_project/mids.py
T
Ankit Malik e218aafc26 first commit
2026-06-22 16:03:00 +05:30

75 lines
1.5 KiB
Python

import os
# import pyarrow
import sys
import logging
from datetime import date, timedelta
import polars as pl
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine, URL
import clickhouse_connect
from dotenv import load_dotenv
from log import log
from clickhouse_task.create_table import *
from db_con.connection import *
def MID_TABLE_COV(engine: Engine, target_date: date) -> list[int]:
sql = text("""
SELECT MID FROM OneApp_KelloggsMT.dbo.T_StoreCoverage
WHERE CONVERT(date, CreateDate) = :target_date
UNION
SELECT MID FROM OneApp_KelloggsMT.dbo.T_StoreCoverage
WHERE CONVERT(date, UpdateDate) = :target_date
""")
log.info(f"Collecting MIDs for: {target_date}")
with engine.connect() as conn:
result = conn.execute(sql, {"target_date": target_date})
mids = [row[0] for row in result.fetchall()]
log.info(f"Found {len(mids):,} MIDs")
return mids
def MID_TABLE_COV1(
engine: Engine,
target_date: date,
) -> pl.DataFrame:
query = f"""
SELECT
EmpId,
CAST(VisitDate AS DATE) AS VisitDate
FROM OneApp_KelloggsMT.dbo.T_OQAD
WHERE CAST(CreateDate AS DATE) = '{target_date}'
UNION
SELECT
EmpId,
CAST(VisitDate AS DATE) AS VisitDate
FROM OneApp_KelloggsMT.dbo.T_OQAD
WHERE CAST(UpdateDate AS DATE) = '{target_date}'
"""
return pl.read_database(
query=query,
connection=engine,
)