from azure.identity import AzureCliCredential
from sqlalchemy import create_engine
from sqlalchemy.dialects import registry
from sqlalchemy.engine import URL


# Dynamically register the KQLDialect plugin
registry.register("mskql.pyodbc", "kqlalchemy.pyodbc", "KQLDialect_pyodbc")

server = "help"
database = "Samples"
conn_str = f"Driver={{ODBC Driver 17 for SQL Server}};Server={server}.kusto.windows.net;Database={database}"
connection_url = URL.create("mskql+pyodbc", query={"odbc_connect": conn_str, "autocommit": "True"})
engine = create_engine(connection_url, azure_credential=AzureCliCredential())

from azure.identity import AzureCliCredential
# Importing kqlalchemy loads the dialect plugin.
from kqlalchemy import kusto_engine, kusto_table
import pandas as pd
from sqlalchemy.orm import Session

engine = kusto_engine("help", "Samples", AzureCliCredential())
storm_events = kusto_table("StormEvents", engine) # performs reflection
query = storm_events.select().limit(10)

with Session(engine, autocommit=True) as session:
    df = pd.read_sql(
        query.compile(engine, compile_kwargs={"literal_binds": True}),
        session.bind,
    )
    print(df[["StartTime", "EventType", "State", "DamageProperty"]].head(5))