azure-log-analytics-data-export.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "357a3e78-a2de-4b57-9cbf-aca900106188",
   "metadata": {},
   "source": [
    "# Azure Log Analytics Data Export"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42f79204-c39d-44d5-a49e-a9342e84784a",
   "metadata": {},
   "source": [
    "This notebook demonstrates the export of data (10M+ records per hour) from Azure Log Analytics to Blob Storage via Python SDKs. In production, this can be deployed via Azure Functions. In testing, 50M records with 10 columns was successfully exported in approximately 1 hour using an Azure Function App (Consumption or Serverless hosting plan). \n",
    "\n",
    "Note: Requires the following Azure resources:\n",
    "\n",
    "- Log Analytics Workspace (data source)\n",
    "- Storage Account (Container, Queue, and Table)\n",
    "  - Container (export destination)\n",
    "  - Queue (split query jobs/messages)\n",
    "  - Table (logging)  \n",
    "\n",
    "Note: Your authentication method (AAD/Entra account or service principal) requires the following roles:\n",
    "\n",
    "- Monitoring Metrics Publisher (Ingest to Log Analytics)\n",
    "- Log Analytics Contributor (Query Log Analytics)\n",
    "- Storage Queue Data Contributor (Storage Queue Send/Get/Delete)\n",
    "- Storage Queue Data Message Processor (Storage Queue Trigger for Azure Function)\n",
    "- Storage Blob Data Contributor (Upload to Blob Storage)\n",
    "- Storage Table Data Contributor (Write to Storage Table for logging)\n",
    "\n",
    "Inputs and Outputs:\n",
    "- <b>Input:</b> table(s), columns, and date range \n",
    "- <b>Output:</b> json, csv, or parquet files \n",
    " \n",
    "Summary:\n",
    "1. <b>Generate Test Data:</b> creates and ingests test data \n",
    "2. <b>Split Query and Send to Queue:</b> divides request into smaller queries/jobs and sends to storage queue \n",
    "3. <b>Process Queue:</b> runs jobs from the storage queue and saves to storage account"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8b6023b-89e4-4fca-9dae-5562ad94a280",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cb22b8d-138f-4e10-affd-c2dcfabc8d7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import hashlib\n",
    "import importlib.util\n",
    "import json\n",
    "import logging\n",
    "import math\n",
    "import os\n",
    "import random\n",
    "import string\n",
    "import time\n",
    "import uuid\n",
    "from io import BytesIO, StringIO\n",
    "from pathlib import Path\n",
    "\n",
    "import pandas as pd\n",
    "from azure.data.tables import TableClient, UpdateMode\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.monitor.ingestion import LogsIngestionClient\n",
    "from azure.monitor.query import LogsQueryClient, LogsQueryStatus\n",
    "from azure.storage.blob import ContainerClient\n",
    "from azure.storage.queue import QueueClient, QueueMessage\n",
    "\n",
    "# pyarrow required for pandas parquet output, check without importing\n",
    "check_packages_installed = [\"pyarrow\"]\n",
    "for each_package in check_packages_installed:\n",
    "    if not importlib.util.find_spec(each_package):\n",
    "        raise Exception(f\"Failed to import package: {each_package}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "591e6d87-77bc-46e4-a44f-0caaccf1bd35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------------------------\n",
    "# notebook settings\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "# increase column display for longer filenames\n",
    "pd.set_option(\"max_colwidth\", None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b81f8537-5b0a-4dc6-ab14-dc50344c712b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------------------------\n",
    "# logging\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "# output directory\n",
    "output_log_dir = \"logs\"\n",
    "if not Path(output_log_dir).is_dir():\n",
    "    Path(output_log_dir).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# logging configuration\n",
    "logging_timestamp = pd.Timestamp.today(\"UTC\").strftime(\"%Y-%m-%d %H:%M:%S.%f\")\n",
    "logging_timestamp = (\n",
    "    logging_timestamp.replace(\"-\", \"\").replace(\":\", \"\").replace(\".\", \"\").replace(\" \", \"\")\n",
    ")\n",
    "logging.basicConfig(\n",
    "    filename=f\"./logs/log-analytics-data-export-{logging_timestamp}.log\",\n",
    "    format=\"%(asctime)s %(levelname)s %(message)s\",\n",
    "    filemode=\"w\",\n",
    ")\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "514b5735-0545-4028-bf78-8a972b39112f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------------------------\n",
    "# log analytics ingest\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "def break_up_ingest_requests(\n",
    "    start_datetime: str,\n",
    "    time_delta_seconds: float,\n",
    "    number_of_rows: int,\n",
    "    max_rows_per_request: int,\n",
    ") -> pd.DataFrame:\n",
    "    number_of_loops = math.ceil(number_of_rows / max_rows_per_request)\n",
    "    next_start_datetime = pd.to_datetime(start_datetime)\n",
    "    rows_to_generate = number_of_rows\n",
    "    ingest_requests = []\n",
    "    for _ in range(number_of_loops):\n",
    "        # start datetimes\n",
    "        each_ingest_request = {}\n",
    "        each_next_start_datetime = next_start_datetime.strftime(\"%Y-%m-%d %H:%M:%S.%f\")\n",
    "        each_ingest_request[\"start_datetime\"] = each_next_start_datetime\n",
    "        # determine number of rows for each request\n",
    "        if rows_to_generate < max_rows_per_request:\n",
    "            request_number_of_rows = rows_to_generate\n",
    "        else:\n",
    "            request_number_of_rows = max_rows_per_request\n",
    "        each_ingest_request[\"number_of_rows\"] = request_number_of_rows\n",
    "        ingest_requests.append(each_ingest_request)\n",
    "        # update number of rows and start datetime for next request\n",
    "        rows_to_generate -= request_number_of_rows\n",
    "        next_start_datetime += pd.to_timedelta(\n",
    "            request_number_of_rows * time_delta_seconds, unit=\"s\"\n",
    "        )\n",
    "    ingest_requests_df = pd.DataFrame(ingest_requests)\n",
    "    return ingest_requests_df\n",
    "\n",
    "\n",
    "def generate_test_data(\n",
    "    start_date: str,\n",
    "    timedelta_seconds: int,\n",
    "    number_of_rows: int,\n",
    "    number_of_columns: int,\n",
    "    random_length: int = 10,\n",
    ") -> pd.DataFrame:\n",
    "    # create dataframe\n",
    "    start_datetime = pd.to_datetime(start_date)\n",
    "    timedelta = pd.Series(range(number_of_rows)) * pd.to_timedelta(\n",
    "        f\"{timedelta_seconds}s\"\n",
    "    )\n",
    "    fake_time_column = start_datetime + timedelta\n",
    "    fake_data_df = pd.DataFrame(\n",
    "        {\n",
    "            \"TimeGenerated\": fake_time_column,\n",
    "        }\n",
    "    )\n",
    "    for each_index in range(1, number_of_columns):\n",
    "        each_column_name = f\"DataColumn{each_index}\"\n",
    "        each_column_value = \"\".join(\n",
    "            random.choice(string.ascii_lowercase) for i in range(random_length)\n",
    "        )\n",
    "        fake_data_df[each_column_name] = each_column_value\n",
    "    # convert datetime to string column to avoid issues in log analytics\n",
    "    time_generated = fake_data_df[\"TimeGenerated\"].dt.strftime(\"%Y-%m-%d %H:%M:%S.%f\")\n",
    "    fake_data_df[\"TimeGenerated\"] = time_generated\n",
    "    # status\n",
    "    logging.info(\"Data Shape: %s\", fake_data_df.shape)\n",
    "    size_calculation = fake_data_df.memory_usage().sum() / 1_000_000\n",
    "    logging.info(\"Size: %s MBs\", size_calculation)\n",
    "    get_first_datetime = fake_data_df[\"TimeGenerated\"].iloc[0]\n",
    "    get_last_datetime = fake_data_df[\"TimeGenerated\"].iloc[-1]\n",
    "    logging.info(\"First Datetime: %s\", get_first_datetime)\n",
    "    logging.info(\"Last Datetime: %s\", get_last_datetime)\n",
    "    return fake_data_df\n",
    "\n",
    "\n",
    "def log_analytics_ingest(\n",
    "    fake_data_df: pd.DataFrame,\n",
    "    ingest_client: LogsIngestionClient,\n",
    "    rule_id: str,\n",
    "    stream_name: str,\n",
    ") -> int:\n",
    "    try:\n",
    "        # convert to json\n",
    "        body = json.loads(fake_data_df.to_json(orient=\"records\", date_format=\"iso\"))\n",
    "        # send to log analytics\n",
    "        ingest_client.upload(rule_id=rule_id, stream_name=stream_name, logs=body)\n",
    "        logging.info(\"Send Successful\")\n",
    "        # return count of rows\n",
    "        return fake_data_df.shape[0]\n",
    "    except Exception as e:\n",
    "        logging.error(\"Error sending to log analytics, will skip: %s\", e)\n",
    "        return 0\n",
    "\n",
    "\n",
    "def generate_and_ingest_test_data(\n",
    "    credential: DefaultAzureCredential,\n",
    "    endpoint: str,\n",
    "    rule_id: str,\n",
    "    stream_name: str,\n",
    "    storage_table_url: str,\n",
    "    storage_table_ingest_name: str,\n",
    "    start_date: str,\n",
    "    timedelta_seconds: float,\n",
    "    number_of_rows: int,\n",
    "    number_of_columns: int = 10,\n",
    "    max_rows_per_request=5_000_000,\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    Generates test/fake data and ingests in Log Analytics\n",
    "        note: credential requires Log Analytics Contributor and Monitor Publisher roles\n",
    "        note: 10M rows with 10 columns takes about 15-20 minutes\n",
    "    Log Analytics Data Collection Endpoint and Rule setup:\n",
    "        1. azure portal -> monitor -> create data collection endpoint\n",
    "        2. azure portal -> log analytics -> table -> create new custom table in log analytics\n",
    "        3. create data collection rule and add publisher role permissions\n",
    "        reference: https://learn.microsoft.com/en-us/azure/azure-monitor/logs/tutorial-logs-ingestion-portal\n",
    "    Args:\n",
    "        credential: DefaultAzureCredential\n",
    "        endpoint: log analytics endpoint url\n",
    "            format: \"https://{name}-XXXX.eastus-1.ingest.monitor.azure.com\"\n",
    "        rule_id: required log analytics ingestion param\n",
    "            format: \"dcr-XXXXXXXXXXXXXXXXXXXXXXXXXXXX\"\n",
    "        stream_name: required log analytics ingestion param\n",
    "            format: \"Custom-{tablename}\"\n",
    "        storage_table_url: url for storage table\n",
    "            format: \"https://{storage_account_name}.table.core.windows.net/\"\n",
    "        storage_table_ingest_name: name of storage table for ingest logs\n",
    "        start_date: date to insert fake data\n",
    "            format: YYYY-MM-DD HH:MM:SS\n",
    "            note: can only ingest dates up to 2 days in the past and 1 day into the future\n",
    "            reference: https://learn.microsoft.com/en-us/azure/azure-monitor/logs/log-standard-columns\n",
    "        timedelta_seconds: time between each fake data row\n",
    "        number_of_rows: total number of rows to generate\n",
    "        number_of_columns: total number of columns to generate\n",
    "            note: for new columns, you need to update the schema before ingestion\n",
    "            1. azure portal -> log analytics -> settings - tables -> ... -> edit schema\n",
    "            2. azure portal -> data collection rules -> export template -> deploy -> edit\n",
    "        max_rows_per_request: limit on number of rows to generate for each ingest\n",
    "            note: lower this if running out memory\n",
    "            note: 5M rows with 10 columns requires about 4-8 GB of RAM\n",
    "    Returns:\n",
    "        dict with results summary\n",
    "    \"\"\"\n",
    "    time_start = time.time()\n",
    "    # input validation\n",
    "    given_timestamp = pd.to_datetime(start_date)\n",
    "    current_datetime = pd.to_datetime(\"today\")\n",
    "    check_start_range = current_datetime - pd.to_timedelta(\"2D\")\n",
    "    check_end_range = current_datetime + pd.to_timedelta(\"1D\")\n",
    "    if not check_start_range <= given_timestamp <= check_end_range:\n",
    "        logging.warning(\"Warning: Date given is outside allowed ingestion range\")\n",
    "        logging.warning(\"Note: Log Analytics will use ingest time as TimeGenerated\")\n",
    "        valid_ingest_datetime_range = False\n",
    "    else:\n",
    "        valid_ingest_datetime_range = True\n",
    "    if number_of_rows < 2 or number_of_columns < 2:\n",
    "        raise Exception(\"invalid row and/or column numbers\")\n",
    "    # log analytics ingest connection\n",
    "    ingest_client = LogsIngestionClient(endpoint, credential)\n",
    "    # storage table connection for logging\n",
    "    # note: requires Storage Table Data Contributor role\n",
    "    table_client = TableClient(\n",
    "        storage_table_url, storage_table_ingest_name, credential=credential\n",
    "    )\n",
    "    # break up ingests\n",
    "    ingest_requests_df = break_up_ingest_requests(\n",
    "        start_date, timedelta_seconds, number_of_rows, max_rows_per_request\n",
    "    )\n",
    "    number_of_ingests = len(ingest_requests_df)\n",
    "    # loop through requests\n",
    "    successfull_rows_sent = 0\n",
    "    for each_row in ingest_requests_df.itertuples():\n",
    "        each_index = each_row.Index + 1\n",
    "        each_request_start_time = time.time()\n",
    "        each_start_datetime = each_row.start_datetime\n",
    "        each_number_of_rows = each_row.number_of_rows\n",
    "        # generate fake data\n",
    "        logging.info(\"Generating Request %s of %s...\", each_index, number_of_ingests)\n",
    "        try:\n",
    "            each_fake_data_df = generate_test_data(\n",
    "                each_start_datetime,\n",
    "                timedelta_seconds,\n",
    "                each_number_of_rows,\n",
    "                number_of_columns,\n",
    "            )\n",
    "        except Exception as e:\n",
    "            logging.error(\"Unable to generate test data: %s\", e)\n",
    "            continue\n",
    "        # send to log analytics\n",
    "        logging.info(\"Sending to Log Analytics...\")\n",
    "        each_rows_ingested = log_analytics_ingest(\n",
    "            each_fake_data_df,\n",
    "            ingest_client,\n",
    "            rule_id,\n",
    "            stream_name,\n",
    "        )\n",
    "        successfull_rows_sent += each_rows_ingested\n",
    "        runtime_calculation = round(time.time() - each_request_start_time, 1)\n",
    "        logging.info(\"Runtime: %s seconds\", runtime_calculation)\n",
    "    # status check\n",
    "    if successfull_rows_sent == 0:\n",
    "        status = \"Failed\"\n",
    "    elif successfull_rows_sent == number_of_rows:\n",
    "        status = \"Success\"\n",
    "    else:\n",
    "        status = \"Partial\"\n",
    "    # create partition key and row key\n",
    "    ingest_uuid = str(uuid.uuid4())\n",
    "    first_datetime = pd.to_datetime(start_date).strftime(\"%Y-%m-%d %H:%M:%S.%f\")\n",
    "    last_datetime = each_fake_data_df[\"TimeGenerated\"].iloc[-1]\n",
    "    row_key = f\"{ingest_uuid}__{status}__\"\n",
    "    row_key += f\"{first_datetime}__{last_datetime}__{timedelta_seconds}__\"\n",
    "    row_key += f\"{number_of_columns}__{number_of_rows}__{successfull_rows_sent}\"\n",
    "    unique_row_sha256_hash = hashlib.sha256(row_key.encode()).hexdigest()\n",
    "    # response and logging to table storage\n",
    "    runtime = round(time.time() - time_start, 1)\n",
    "    time_generated = pd.Timestamp.today(\"UTC\").strftime(\"%Y-%m-%d %H:%M:%S.%f\")\n",
    "    return_message = {\n",
    "        \"PartitionKey\": ingest_uuid,\n",
    "        \"RowKey\": unique_row_sha256_hash,\n",
    "        \"Status\": status,\n",
    "        \"StartDatetime\": first_datetime,\n",
    "        \"EndDatetime\": last_datetime,\n",
    "        \"TimeDeltaSeconds\": timedelta_seconds,\n",
    "        \"NumberColumns\": number_of_columns,\n",
    "        \"RowsGenerated\": number_of_rows,\n",
    "        \"RowsIngested\": successfull_rows_sent,\n",
    "        \"ValidDatetimeRange\": valid_ingest_datetime_range,\n",
    "        \"RuntimeSeconds\": runtime,\n",
    "        \"TimeGenerated\": time_generated,\n",
    "    }\n",
    "    table_client.upsert_entity(return_message, mode=UpdateMode.REPLACE)\n",
    "    return return_message\n",
    "\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# log analytics query\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "def query_log_analytics_request(\n",
    "    workspace_id: str,\n",
    "    log_client: LogsQueryClient,\n",
    "    kql_query: str,\n",
    "    request_wait_seconds: float = 0.05,\n",
    ") -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Makes API query request to log analytics\n",
    "    limits: https://learn.microsoft.com/en-us/azure/azure-monitor/logs/api/timeouts\n",
    "    API query limits:\n",
    "        500,000 rows per request\n",
    "        200 requests per 30 seconds\n",
    "        max query time is 10 min\n",
    "        100MB data max per request\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # query log analytics\n",
    "        response = log_client.query_workspace(\n",
    "            workspace_id=workspace_id,\n",
    "            query=kql_query,\n",
    "            timespan=None,\n",
    "            server_timeout=600,\n",
    "        )\n",
    "        # convert to dataframe\n",
    "        if response.status == LogsQueryStatus.SUCCESS:\n",
    "            table = response.tables[0]\n",
    "            df = pd.DataFrame(data=table.rows, columns=table.columns)\n",
    "            return df\n",
    "        elif response.status == LogsQueryStatus.PARTIAL:\n",
    "            raise Exception(\n",
    "                f\"Unsuccessful Request, Response Status: {response.status} {response.partial_error}\"\n",
    "            )\n",
    "        else:\n",
    "            raise Exception(\n",
    "                f\"Unsuccessful Request, Response Status: {response.status} {response}\"\n",
    "            )\n",
    "    except Exception as e:\n",
    "        raise Exception(\"Failed Log Analytics Request\") from e\n",
    "    finally:\n",
    "        time.sleep(request_wait_seconds)\n",
    "\n",
    "\n",
    "def query_log_analytics_connection_request(\n",
    "    credential: DefaultAzureCredential, workspace_id: str, kql_query: str\n",
    ") -> pd.DataFrame:\n",
    "    # log analytics connection\n",
    "    # note: need to add Log Analytics Contributor and Monitor Publisher role\n",
    "    log_client = LogsQueryClient(credential)\n",
    "    # submit query request\n",
    "    result_df = query_log_analytics_request(workspace_id, log_client, kql_query)\n",
    "    return result_df\n",
    "\n",
    "\n",
    "def query_log_analytics_get_table_columns(\n",
    "    table_names_and_columns: dict,\n",
    "    workspace_id: str,\n",
    "    log_client: LogsQueryClient,\n",
    ") -> dict:\n",
    "    output = {}\n",
    "    for each_table, each_columns in table_names_and_columns.items():\n",
    "        # column names provided\n",
    "        if each_columns:\n",
    "            each_columns_fix = each_columns\n",
    "            if \"TimeGenerated\" not in each_columns:\n",
    "                each_columns_fix = [\"TimeGenerated\"] + each_columns\n",
    "            output[each_table] = each_columns_fix\n",
    "        # if no column names provided, query log analytics for all column names\n",
    "        else:\n",
    "            logging.info(\"Getting columns names for %s\", each_table)\n",
    "            each_kql_query = f\"\"\"\n",
    "            let TABLE_NAME = \"{each_table}\";\n",
    "            table(TABLE_NAME)\n",
    "            | project-away TenantId, Type, _ResourceId\n",
    "            | take 1\n",
    "            \"\"\"\n",
    "            each_df = query_log_analytics_request(\n",
    "                workspace_id, log_client, each_kql_query\n",
    "            )\n",
    "            each_columns_fix = list(each_df.columns)\n",
    "            each_columns_fix.remove(\"TimeGenerated\")\n",
    "            each_columns_fix = [\"TimeGenerated\"] + each_columns_fix\n",
    "            logging.info(\"Columns Detected: %s\", each_columns_fix)\n",
    "            output[each_table] = each_columns_fix\n",
    "    if len(output) == 0:\n",
    "        raise Exception(\"No valid table names\")\n",
    "    return output\n",
    "\n",
    "\n",
    "def break_up_initial_date_range(\n",
    "    table_name: str, start_datetime: str, end_datetime: str, freq: str\n",
    ") -> pd.DataFrame:\n",
    "    date_range = pd.date_range(start=start_datetime, end=end_datetime, freq=freq)\n",
    "    if date_range[-1] != pd.to_datetime(end_datetime):\n",
    "        date_range = date_range.union(pd.to_datetime([end_datetime]))\n",
    "    date_ranges = [each.strftime(\"%Y-%m-%d %H:%M:%S.%f\") for each in date_range.to_list()]\n",
    "    time_pairs = [\n",
    "        (date_ranges[i], date_ranges[i + 1]) for i in range(len(date_ranges) - 1)\n",
    "    ]\n",
    "    # convert to dataframe\n",
    "    df_time_pairs = pd.DataFrame(time_pairs, columns=[\"start_date\", \"end_date\"])\n",
    "    df_time_pairs.insert(loc=0, column=\"table\", value=[table_name] * len(df_time_pairs))\n",
    "    return df_time_pairs\n",
    "\n",
    "\n",
    "def break_up_initial_query_time_freq(\n",
    "    table_names: list[str], start_datetime: str, end_datetime: str, freq: str\n",
    ") -> pd.DataFrame:\n",
    "    results = []\n",
    "    # break up by table names\n",
    "    for each_table_name in table_names:\n",
    "        # break up date ranges by day\n",
    "        each_df = break_up_initial_date_range(\n",
    "            each_table_name, start_datetime, end_datetime, freq\n",
    "        )\n",
    "        results.append(each_df)\n",
    "    df_results = pd.concat(results)\n",
    "    return df_results\n",
    "\n",
    "\n",
    "def query_log_analytics_get_time_ranges(\n",
    "    workspace_id: str,\n",
    "    log_client: LogsQueryClient,\n",
    "    table_name: str,\n",
    "    start_datetime: str,\n",
    "    end_datetime: str,\n",
    "    query_row_limit: int,\n",
    ") -> pd.DataFrame:\n",
    "    # converted KQL output to string columns to avoid datetime digits getting truncated\n",
    "    kql_query = f\"\"\"\n",
    "    let TABLE_NAME = \"{table_name}\";\n",
    "    let START_DATETIME = datetime({start_datetime});\n",
    "    let END_DATETIME = datetime({end_datetime});\n",
    "    let QUERY_ROW_LIMIT = {query_row_limit};\n",
    "    let table_datetime_filtered = table(TABLE_NAME)\n",
    "    | project TimeGenerated\n",
    "    | where (TimeGenerated >= START_DATETIME) and (TimeGenerated < END_DATETIME);\n",
    "    let table_size = toscalar(\n",
    "    table_datetime_filtered\n",
    "    | count);\n",
    "    let time_splits = table_datetime_filtered\n",
    "    | order by TimeGenerated asc\n",
    "    | extend row_index = row_number()\n",
    "    | where row_index == 1 or row_index % (QUERY_ROW_LIMIT) == 0 or row_index == table_size;\n",
    "    let time_pairs = time_splits\n",
    "    | project StartTime = TimeGenerated\n",
    "    | extend EndTime = next(StartTime)\n",
    "    | where isnotnull(EndTime)\n",
    "    | extend StartTime = tostring(StartTime), EndTime = tostring(EndTime);\n",
    "    time_pairs\n",
    "    \"\"\"\n",
    "    logging.info(\"Splitting %s: %s-%s\", table_name, start_datetime, end_datetime)\n",
    "    # query log analytics and get time ranges\n",
    "    df = query_log_analytics_request(workspace_id, log_client, kql_query)\n",
    "    # no results\n",
    "    if df.shape[0] == 0:\n",
    "        return pd.DataFrame()\n",
    "    # datetime fix for events on final datetime\n",
    "    # using copy and .loc to prevent chaining warning\n",
    "    df_copy = df.copy()\n",
    "    final_endtime = df_copy[\"EndTime\"].tail(1).item()\n",
    "    new_final_endtime = str(pd.to_datetime(final_endtime) + pd.to_timedelta(\"0.0000001s\"))\n",
    "    new_final_endtime_fix_format = new_final_endtime.replace(\" \", \"T\")\n",
    "    new_final_endtime_fix_format = new_final_endtime_fix_format.replace(\"00+00:00\", \"Z\")\n",
    "    df_copy.loc[df_copy.index[-1], \"EndTime\"] = new_final_endtime_fix_format\n",
    "    return df_copy\n",
    "\n",
    "\n",
    "def query_log_analytics_get_table_count(\n",
    "    workspace_id: str,\n",
    "    log_client: LogsQueryClient,\n",
    "    table_name: str,\n",
    "    start_datetime: str,\n",
    "    end_datetime: str,\n",
    ") -> int:\n",
    "    kql_query = f\"\"\"\n",
    "    let TABLE_NAME = \"{table_name}\";\n",
    "    let START_DATETIME = datetime({start_datetime});\n",
    "    let END_DATETIME = datetime({end_datetime});\n",
    "    table(TABLE_NAME)\n",
    "    | project TimeGenerated\n",
    "    | where (TimeGenerated >= START_DATETIME) and (TimeGenerated < END_DATETIME)\n",
    "    | count\n",
    "    \"\"\"\n",
    "    df = query_log_analytics_request(workspace_id, log_client, kql_query)\n",
    "    return df.values[0][0]\n",
    "\n",
    "\n",
    "def query_log_analytics_add_table_row_counts(\n",
    "    input_df: pd.DataFrame,\n",
    "    workspace_id: str,\n",
    "    log_client: LogsQueryClient,\n",
    "    table_name: str,\n",
    ") -> pd.DataFrame:\n",
    "    # add row counts\n",
    "    results = []\n",
    "    for each_row in input_df.itertuples():\n",
    "        each_starttime = each_row.StartTime\n",
    "        each_endtime = each_row.EndTime\n",
    "        each_count = query_log_analytics_get_table_count(\n",
    "            workspace_id, log_client, table_name, each_starttime, each_endtime\n",
    "        )\n",
    "        results.append(each_count)\n",
    "    input_df[\"Count\"] = results\n",
    "    return input_df\n",
    "\n",
    "\n",
    "def query_log_analytics_split_query_rows(\n",
    "    workspace_id: str,\n",
    "    log_client: LogsQueryClient,\n",
    "    table_name: str,\n",
    "    start_datetime: str,\n",
    "    end_datetime: str,\n",
    "    query_row_limit: int,\n",
    "    query_row_limit_correction: int,\n",
    ") -> pd.DataFrame:\n",
    "    # fix for large number of events at same datetime\n",
    "    query_row_limit_fix = query_row_limit - query_row_limit_correction\n",
    "    # get time ranges\n",
    "    results_df = query_log_analytics_get_time_ranges(\n",
    "        workspace_id,\n",
    "        log_client,\n",
    "        table_name,\n",
    "        start_datetime,\n",
    "        end_datetime,\n",
    "        query_row_limit_fix,\n",
    "    )\n",
    "    # empty results\n",
    "    if results_df.shape[0] == 0:\n",
    "        return pd.DataFrame()\n",
    "    # add row counts column\n",
    "    results_df = query_log_analytics_add_table_row_counts(\n",
    "        results_df, workspace_id, log_client, table_name\n",
    "    )\n",
    "    # warning if query limit exceeded, change limits and try again\n",
    "    if results_df.Count.gt(query_row_limit).any():\n",
    "        raise Exception(f\"Sub-Query exceeds query row limit, {list(results_df.Count)}\")\n",
    "    # add table name column\n",
    "    results_df.insert(loc=0, column=\"Table\", value=[table_name] * len(results_df))\n",
    "    return results_df\n",
    "\n",
    "\n",
    "def query_log_analytics_split_query_rows_loop(\n",
    "    df_queries: pd.DataFrame,\n",
    "    workspace_id: str,\n",
    "    log_client: LogsQueryClient,\n",
    "    query_row_limit: int,\n",
    "    query_row_limit_correction: int,\n",
    ") -> pd.DataFrame:\n",
    "    logging.info(\"Querying Log Analytics to Split Query...\")\n",
    "    query_results = []\n",
    "    for each_query in df_queries.itertuples():\n",
    "        each_table_name = each_query.table\n",
    "        each_start_datetime = each_query.start_date\n",
    "        each_end_datetime = each_query.end_date\n",
    "        each_results_df = query_log_analytics_split_query_rows(\n",
    "            workspace_id,\n",
    "            log_client,\n",
    "            each_table_name,\n",
    "            each_start_datetime,\n",
    "            each_end_datetime,\n",
    "            query_row_limit,\n",
    "            query_row_limit_correction,\n",
    "        )\n",
    "        query_results.append(each_results_df)\n",
    "        # each status\n",
    "        each_status = f\"Completed {each_table_name}: \"\n",
    "        each_status += f\"{each_start_datetime}-{each_end_datetime} \"\n",
    "        each_status += f\"-> {each_results_df.shape[0]} Queries\"\n",
    "        logging.info(each_status)\n",
    "    # combine all results\n",
    "    results_df = pd.concat(query_results)\n",
    "    return results_df\n",
    "\n",
    "\n",
    "def process_query_results_df(\n",
    "    query_results_df: pd.DataFrame,\n",
    "    query_uuid: str,\n",
    "    table_names_and_columns: dict,\n",
    "    subscription_id: str,\n",
    "    resource_group: str,\n",
    "    worksapce_name: str,\n",
    "    workspace_id: str,\n",
    "    storage_blob_url: str,\n",
    "    storage_blob_name: str,\n",
    "    storage_blob_output: str,\n",
    "    storage_table_url: str,\n",
    "    storage_table_name: str,\n",
    ") -> list[dict]:\n",
    "    # add column names\n",
    "    column_names = query_results_df[\"Table\"].apply(lambda x: table_names_and_columns[x])\n",
    "    query_results_df.insert(loc=1, column=\"Columns\", value=column_names)\n",
    "    # add azure property columns\n",
    "    query_results_df.insert(loc=0, column=\"QueryUUID\", value=query_uuid)\n",
    "    query_results_df.insert(loc=5, column=\"Subscription\", value=subscription_id)\n",
    "    query_results_df.insert(loc=6, column=\"ResourceGroup\", value=resource_group)\n",
    "    query_results_df.insert(loc=7, column=\"LogAnalyticsWorkspace\", value=worksapce_name)\n",
    "    query_results_df.insert(loc=8, column=\"LogAnalyticsWorkspaceId\", value=workspace_id)\n",
    "    query_results_df.insert(loc=9, column=\"StorageBlobURL\", value=storage_blob_url)\n",
    "    query_results_df.insert(loc=10, column=\"StorageContainer\", value=storage_blob_name)\n",
    "    query_results_df.insert(loc=11, column=\"OutputFormat\", value=storage_blob_output)\n",
    "    query_results_df.insert(loc=12, column=\"StorageTableURL\", value=storage_table_url)\n",
    "    query_results_df.insert(loc=13, column=\"StorageTableName\", value=storage_table_name)\n",
    "    # rename columns\n",
    "    query_results_df_rename = query_results_df.rename(\n",
    "        columns={\"StartTime\": \"StartDatetime\", \"EndTime\": \"EndDatetime\"}\n",
    "    )\n",
    "    # convert to dictionary\n",
    "    results = query_results_df_rename.to_dict(orient=\"records\")\n",
    "    return results\n",
    "\n",
    "\n",
    "def query_log_analytics_send_to_queue(\n",
    "    query_uuid: str,\n",
    "    credential: DefaultAzureCredential,\n",
    "    subscription_id: str,\n",
    "    resource_group: str,\n",
    "    worksapce_name: str,\n",
    "    workspace_id: str,\n",
    "    storage_queue_url: str,\n",
    "    storage_queue_name: str,\n",
    "    storage_blob_url: str,\n",
    "    storage_blob_container: str,\n",
    "    storage_table_url: str,\n",
    "    storage_table_query_name: str,\n",
    "    storage_table_process_name: str,\n",
    "    table_names_and_columns: dict,\n",
    "    start_datetime: str,\n",
    "    end_datetime: str,\n",
    "    query_row_limit: int = 250_000,\n",
    "    query_row_limit_correction: int = 1_000,\n",
    "    break_up_query_freq=\"4h\",\n",
    "    storage_blob_output_format: str = \"JSONL\",\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    Splits query date range into smaller queries and sends to storage queue\n",
    "        note: credential requires Log Analytics, Storage Queue, and Table Storage Contributor roles\n",
    "        note: date range is processed as [start_datetime, end_datetime)\n",
    "    Args:\n",
    "        query_uuid: uuid for full query\n",
    "            format: \"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX\"\n",
    "        credential: azure default credential object\n",
    "        subscription_id: azure subscription id\n",
    "            format: \"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX\"\n",
    "        resource_group: azure resource group\n",
    "        workspace_name: name of log analytics workspace\n",
    "        workspace_id: log analytics workspace id\n",
    "            format: \"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX\"\n",
    "        storage_queue_url: storage account queue url\n",
    "            format: \"https://{storage_account_name}.queue.core.windows.net/\"\n",
    "        storage_queue_name: name of storage queue\n",
    "        storage_blob_url: storage blob account url to save output\n",
    "            format: https://{storage_account_name}.blob.core.windows.net/\"\n",
    "        storage_blob_container: name of container in storage account to save output\n",
    "        storage_table_url: storage table url\n",
    "            format: \"https://{storage_account_name}.table.core.windows.net/\"\n",
    "        storage_table_query_name: name of storage table for query logs\n",
    "        storage_table_process_name: name of storage table for process logs\n",
    "        table_names_and_columns: dictionary of table names with columns to project\n",
    "            note: blank column list will detect and use all columns\n",
    "            format:  {\"table_name\" : [\"column_1\", \"column_2\", ... ], ... }\n",
    "        start_datetime: starting datetime, inclusive\n",
    "            format: YYYY-MM-DD HH:MM:SS\n",
    "        start_datetime: ending datetime, exclusive\n",
    "            format: YYYY-MM-DD HH:MM:SS\n",
    "        query_row_limit: max number of rows for each follow-up query/message\n",
    "        query_row_limit_correction: correction factor in case of overlapping data\n",
    "        break_up_query_freq: limit on query datetime range to prevent crashes\n",
    "            note:for  more than 10M rows per hour, use 4 hours or less\n",
    "        storage_blob_output_format: output file format, options = \"JSONL\", \"CSV\", \"PARQUET\"\n",
    "            note: JSONL is json line delimited\n",
    "    Return\n",
    "        dict of results summary\n",
    "    \"\"\"\n",
    "    start_time = time.time()\n",
    "    # input validation\n",
    "    try:\n",
    "        pd.to_datetime(start_datetime)\n",
    "        pd.to_datetime(end_datetime)\n",
    "    except Exception as e:\n",
    "        raise Exception(\"Invalid Datetime Format\") from e\n",
    "    if storage_blob_output_format not in [\"JSONL\", \"CSV\", \"PARQUET\"]:\n",
    "        raise Exception(f\"Invalid Output file format: {storage_blob_output_format}\")\n",
    "    # status message\n",
    "    logging.info(\"Processing Query...\")\n",
    "    table_names_join = \", \".join(table_names_and_columns.keys())\n",
    "    logging.info(\"Tables: %s\", table_names_join)\n",
    "    logging.info(\"Date Range: %s-%s\", start_datetime, end_datetime)\n",
    "    # log analytics connection\n",
    "    # note: need to add Log Analytics Contributor role\n",
    "    log_client = LogsQueryClient(credential)\n",
    "    # storage queue connection\n",
    "    # note: need to add Storage Queue Data Contributor role\n",
    "    storage_queue_url_and_name = storage_queue_url + storage_queue_name\n",
    "    queue_client = QueueClient.from_queue_url(storage_queue_url_and_name, credential)\n",
    "    # storage table connection for logging\n",
    "    # note: requires Storage Table Data Contributor role\n",
    "    table_client = TableClient(\n",
    "        storage_table_url, storage_table_query_name, credential=credential\n",
    "    )\n",
    "    # process table and column names\n",
    "    table_names_and_columns = query_log_analytics_get_table_columns(\n",
    "        table_names_and_columns, workspace_id, log_client\n",
    "    )\n",
    "    # get expected count of full queries\n",
    "    total_query_results_count_expected = 0\n",
    "    for each_table_name in table_names_and_columns:\n",
    "        each_count = query_log_analytics_get_table_count(\n",
    "            workspace_id, log_client, each_table_name, start_datetime, end_datetime\n",
    "        )\n",
    "        total_query_results_count_expected += each_count\n",
    "    logging.info(\"Total Row Count: %s\", total_query_results_count_expected)\n",
    "    if total_query_results_count_expected > 0:\n",
    "        # break up queries by table and date ranges\n",
    "        table_names = list(table_names_and_columns.keys())\n",
    "        df_queries = break_up_initial_query_time_freq(\n",
    "            table_names, start_datetime, end_datetime, break_up_query_freq\n",
    "        )\n",
    "        # query log analytics, gets datetime splits for row limit\n",
    "        query_results_df = query_log_analytics_split_query_rows_loop(\n",
    "            df_queries,\n",
    "            workspace_id,\n",
    "            log_client,\n",
    "            query_row_limit,\n",
    "            query_row_limit_correction,\n",
    "        )\n",
    "        if not query_results_df.empty:\n",
    "            # confirm count of split queries\n",
    "            total_query_results_count = query_results_df[\"Count\"].sum()\n",
    "            logging.info(\"Split Queries Total Row Count: %s\", total_query_results_count)\n",
    "            if total_query_results_count != total_query_results_count_expected:\n",
    "                raise Exception(\"Error: Row Count Mismatch\")\n",
    "            # process results, add columns, and convert to list of dicts\n",
    "            results = process_query_results_df(\n",
    "                query_results_df,\n",
    "                query_uuid,\n",
    "                table_names_and_columns,\n",
    "                subscription_id,\n",
    "                resource_group,\n",
    "                worksapce_name,\n",
    "                workspace_id,\n",
    "                storage_blob_url,\n",
    "                storage_blob_container,\n",
    "                storage_blob_output_format,\n",
    "                storage_table_url,\n",
    "                storage_table_process_name,\n",
    "            )\n",
    "            number_of_results = len(results)\n",
    "            # send to queue\n",
    "            successful_sends = 0\n",
    "            get_queue_properties = queue_client.get_queue_properties()\n",
    "            logging.info(\"Initial Queue Status: %s\", get_queue_properties)\n",
    "            for each_msg in results:\n",
    "                each_result = send_message_to_queue(queue_client, each_msg)\n",
    "                if each_result == \"Success\":\n",
    "                    successful_sends += 1\n",
    "            logging.info(\"Messages Successfully Sent to Queue: %s\", successful_sends)\n",
    "            get_queue_properties = queue_client.get_queue_properties()\n",
    "            logging.info(\"Updated Queue Status: %s\", get_queue_properties)\n",
    "            if successful_sends == number_of_results:\n",
    "                status = \"Success\"\n",
    "            else:\n",
    "                status = \"Partial\"\n",
    "        # no results\n",
    "        else:\n",
    "            status = \"Failed\"\n",
    "            total_query_results_count = 0\n",
    "            number_of_results = 0\n",
    "            successful_sends = 0\n",
    "            logging.error(\"Error: No Query Messages Generated, No Query Results\")\n",
    "    # no data during datetime range\n",
    "    else:\n",
    "        status = \"Success\"\n",
    "        total_query_results_count = 0\n",
    "        number_of_results = 0\n",
    "        successful_sends = 0\n",
    "        logging.info(\"No data during %s-%s\", start_datetime, end_datetime)\n",
    "    # create hash for RowKey\n",
    "    row_key = f\"{query_uuid}__{status}__{table_names_join}__\"\n",
    "    row_key += f\"{start_datetime}__{end_datetime}__\"\n",
    "    row_key += f\"{total_query_results_count}__{number_of_results}__{successful_sends}\"\n",
    "    unique_row_sha256_hash = hashlib.sha256(row_key.encode()).hexdigest()\n",
    "    # response and logging to table storage\n",
    "    runtime = round(time.time() - start_time, 1)\n",
    "    time_generated = pd.Timestamp.today(\"UTC\").strftime(\"%Y-%m-%d %H:%M:%S.%f\")\n",
    "    return_message = {\n",
    "        \"PartitionKey\": query_uuid,\n",
    "        \"RowKey\": unique_row_sha256_hash,\n",
    "        \"Status\": status,\n",
    "        \"Tables\": table_names_join,\n",
    "        \"StartDatetime\": start_datetime,\n",
    "        \"EndDatetime\": end_datetime,\n",
    "        \"TotalRowCount\": int(total_query_results_count),\n",
    "        \"MessagesGenerated\": number_of_results,\n",
    "        \"MessagesSentToQueue\": successful_sends,\n",
    "        \"RuntimeSeconds\": runtime,\n",
    "        \"TimeGenerated\": time_generated,\n",
    "    }\n",
    "    table_client.upsert_entity(return_message, mode=UpdateMode.REPLACE)\n",
    "    return return_message\n",
    "\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# storage queue\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "def send_message_to_queue(\n",
    "    queue_client: QueueClient, message: str, request_wait_seconds: float = 0.05\n",
    ") -> str:\n",
    "    try:\n",
    "        queue_client.send_message(json.dumps(message))\n",
    "        return \"Success\"\n",
    "    except Exception as e:\n",
    "        logging.error(\n",
    "            \"Error: Unable to send message to queue, skipped: %s, exception: %s\",\n",
    "            message,\n",
    "            e,\n",
    "        )\n",
    "        return \"Failed\"\n",
    "    finally:\n",
    "        time.sleep(request_wait_seconds)\n",
    "\n",
    "\n",
    "def get_message_from_queue(\n",
    "    queue_client: QueueClient,\n",
    "    message_visibility_timeout_seconds: int,\n",
    "    request_wait_seconds: float = 0.05,\n",
    ") -> QueueMessage | None:\n",
    "    # queue calls have built-in 10x retry policy\n",
    "    # ref: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/storage/azure-storage-queue#optional-configuration\n",
    "    try:\n",
    "        queue_message = queue_client.receive_message(\n",
    "            visibility_timeout=message_visibility_timeout_seconds\n",
    "        )\n",
    "        return queue_message\n",
    "    except Exception as e:\n",
    "        logging.error(\"Request Error: Unable to Get Queue Message, %s\", e)\n",
    "        raise Exception(\"Request Error: Unable to Get Queue Message\") from e\n",
    "    finally:\n",
    "        time.sleep(request_wait_seconds)\n",
    "\n",
    "\n",
    "def delete_message_from_queue(\n",
    "    queue_client: QueueClient, queue_message: QueueMessage\n",
    ") -> None:\n",
    "    try:\n",
    "        queue_client.delete_message(queue_message)\n",
    "        logging.info(\"Successfully Deleted Message from Queue\")\n",
    "    except Exception as e:\n",
    "        logging.error(\"Unable to delete message, %s, %s\", queue_message, e)\n",
    "        raise Exception(f\"Unable to delete message, {queue_message}\") from e\n",
    "\n",
    "\n",
    "def check_if_queue_empty_peek_message(queue_client: QueueClient) -> bool:\n",
    "    try:\n",
    "        peek_messages = queue_client.peek_messages()\n",
    "        if not peek_messages:\n",
    "            return True\n",
    "        return False\n",
    "    except Exception as e:\n",
    "        logging.error(\"Unable to peek at queue messages, %s\", e)\n",
    "        return False\n",
    "\n",
    "\n",
    "def query_message_validation_check(message: dict) -> None:\n",
    "    required_fields = [\n",
    "        \"QueryUUID\",\n",
    "        \"Subscription\",\n",
    "        \"ResourceGroup\",\n",
    "        \"LogAnalyticsWorkspace\",\n",
    "        \"LogAnalyticsWorkspaceId\",\n",
    "        \"StorageQueueURL\",\n",
    "        \"StorageQueueName\",\n",
    "        \"StorageBlobURL\",\n",
    "        \"StorageContainer\",\n",
    "        \"StorageTableURL\",\n",
    "        \"StorageTableQueryName\",\n",
    "        \"StorageTableProcessName\",\n",
    "        \"TableNamesColumns\",\n",
    "        \"StartDatetime\",\n",
    "        \"EndDatetime\",\n",
    "        \"QueryRowLimit\",\n",
    "        \"QueryRowLimitCorrection\",\n",
    "        \"BreakUpQueryFreq\",\n",
    "        \"StorageBlobOutputFormat\",\n",
    "    ]\n",
    "    if not all(each_field in message for each_field in required_fields):\n",
    "        logging.error(\"Invalid message, required fields missing: %s\", message)\n",
    "        raise Exception(f\"Invalid message, required fields missing: {message}\")\n",
    "\n",
    "\n",
    "def process_message_validation_check(message: dict) -> None:\n",
    "    required_fields = [\n",
    "        \"QueryUUID\",\n",
    "        \"Table\",\n",
    "        \"Columns\",\n",
    "        \"StartDatetime\",\n",
    "        \"EndDatetime\",\n",
    "        \"Subscription\",\n",
    "        \"ResourceGroup\",\n",
    "        \"LogAnalyticsWorkspace\",\n",
    "        \"LogAnalyticsWorkspaceId\",\n",
    "        \"StorageBlobURL\",\n",
    "        \"StorageContainer\",\n",
    "        \"OutputFormat\",\n",
    "        \"StorageTableURL\",\n",
    "        \"StorageTableName\",\n",
    "        \"Count\",\n",
    "    ]\n",
    "    if not all(each_field in message for each_field in required_fields):\n",
    "        logging.error(\"Invalid message, required fields missing: %s\", message)\n",
    "        raise Exception(f\"Invalid message, required fields missing: {message}\")\n",
    "\n",
    "\n",
    "def query_log_analytics_get_query_results(\n",
    "    log_client: LogsQueryClient, message: dict\n",
    ") -> pd.DataFrame:\n",
    "    # extract message fields\n",
    "    workspace_id = message[\"LogAnalyticsWorkspaceId\"]\n",
    "    table_name = message[\"Table\"]\n",
    "    column_names = message[\"Columns\"]\n",
    "    start_datetime = message[\"StartDatetime\"]\n",
    "    end_datetime = message[\"EndDatetime\"]\n",
    "    # query log analytics\n",
    "    columns_to_project = \", \".join(column_names)\n",
    "    kql_query = f\"\"\"\n",
    "    let TABLE_NAME = \"{table_name}\";\n",
    "    let START_DATETIME = datetime({start_datetime});\n",
    "    let END_DATETIME = datetime({end_datetime});\n",
    "    table(TABLE_NAME)\n",
    "    | project {columns_to_project}\n",
    "    | where (TimeGenerated >= START_DATETIME) and (TimeGenerated < END_DATETIME)\n",
    "    \"\"\"\n",
    "    df = query_log_analytics_request(workspace_id, log_client, kql_query)\n",
    "    return df\n",
    "\n",
    "\n",
    "def datetime_to_filename_safe(user_input: str) -> str:\n",
    "    # remove characters from timestamp to be filename safe/readable\n",
    "    output = user_input.replace(\"-\", \"\").replace(\":\", \"\").replace(\".\", \"\")\n",
    "    output = output.replace(\"T\", \"\").replace(\"Z\", \"\")\n",
    "    output = output.replace(\" \", \"\")\n",
    "    return output\n",
    "\n",
    "\n",
    "def generate_output_filename_base(\n",
    "    message: str,\n",
    "    output_filename_timestamp: pd.Timestamp,\n",
    ") -> str:\n",
    "    # extract message\n",
    "    table_name = message[\"Table\"]\n",
    "    subscription = message[\"Subscription\"]\n",
    "    resource_group = message[\"ResourceGroup\"]\n",
    "    log_analytics_name = message[\"LogAnalyticsWorkspace\"]\n",
    "    start_datetime = message[\"StartDatetime\"]\n",
    "    end_datetime = message[\"EndDatetime\"]\n",
    "    # datetime conversion via pandas: dt.strftime('%Y-%m-%dT%H:%M:%S.%fZ')\n",
    "    # extract datetime values for filename\n",
    "    extract_year = output_filename_timestamp.strftime(\"%Y\")\n",
    "    extract_month = output_filename_timestamp.strftime(\"%m\")\n",
    "    extract_day = output_filename_timestamp.strftime(\"%d\")\n",
    "    extract_hour = output_filename_timestamp.strftime(\"%H\")\n",
    "    # mimics continuous export from log analytics\n",
    "    # https://learn.microsoft.com/en-us/azure/azure-monitor/logs/logs-data-export\n",
    "    output_filename = f\"{table_name}/\"\n",
    "    output_filename += \"WorkspaceResourceId=/\"\n",
    "    output_filename += f\"subscriptions/{subscription}/\"\n",
    "    output_filename += f\"resourcegroups/{resource_group}/\"\n",
    "    output_filename += \"providers/microsoft.operationalinsights/\"\n",
    "    output_filename += f\"workspaces/{log_analytics_name}/\"\n",
    "    output_filename += f\"y={extract_year}/m={extract_month}/d={extract_day}/\"\n",
    "    output_filename += f\"h={extract_hour}/\"\n",
    "    output_filename += f\"{datetime_to_filename_safe(start_datetime)}-\"\n",
    "    output_filename += f\"{datetime_to_filename_safe(end_datetime)}\"\n",
    "    return output_filename\n",
    "\n",
    "\n",
    "def output_filename_and_format(\n",
    "    results_df: pd.DataFrame, output_format: str, output_filename_base: str\n",
    ") -> tuple[bytes | str]:\n",
    "    # file format\n",
    "    output_filename = output_filename_base\n",
    "    if output_format == \"JSONL\":\n",
    "        output_filename += \".json\"\n",
    "        output_data = results_df.to_json(\n",
    "            orient=\"records\", lines=True, date_format=\"iso\", date_unit=\"ns\"\n",
    "        )\n",
    "    elif output_format == \"CSV\":\n",
    "        output_filename += \".csv\"\n",
    "        output_data = results_df.to_csv(index=False)\n",
    "    elif output_format == \"PARQUET\":\n",
    "        output_filename += \".parquet\"\n",
    "        output_data = results_df.to_parquet(index=False, engine=\"pyarrow\")\n",
    "    return output_filename, output_data\n",
    "\n",
    "\n",
    "def process_queue_message(\n",
    "    log_client: LogsQueryClient,\n",
    "    message: dict,\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Processes individual message: validates, queries log analytics, and saves results to storage account\n",
    "    Args:\n",
    "        log_client: azure log analytics LogsQueryClient object\n",
    "        message: message content dictionary\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    start_time = time.time()\n",
    "    # validate message\n",
    "    process_message_validation_check(message)\n",
    "    logging.info(\"Processing Message: %s\", message)\n",
    "    # query log analytics\n",
    "    query_results_df = query_log_analytics_get_query_results(log_client, message)\n",
    "    logging.info(\"Successfully Downloaded from Log Analytics: %s\", query_results_df.shape)\n",
    "    # confirm count matches\n",
    "    if query_results_df.shape[0] != message[\"Count\"]:\n",
    "        logging.error(\"Row count doesn't match expected value, %s\", message)\n",
    "        raise Exception(f\"Row count doesn't match expected value, {message}\")\n",
    "    # storage blob connection\n",
    "    # note: need to add Storage Blob Data Contributor role\n",
    "    storage_blob_url = message[\"StorageBlobURL\"]\n",
    "    storage_container_name = message[\"StorageContainer\"]\n",
    "    container_client = ContainerClient(\n",
    "        storage_blob_url, storage_container_name, credential\n",
    "    )\n",
    "    # storage table connection for logging\n",
    "    # note: requires Storage Table Data Contributor role\n",
    "    storage_table_url = message[\"StorageTableURL\"]\n",
    "    storage_table_name = message[\"StorageTableName\"]\n",
    "    table_client = TableClient(\n",
    "        storage_table_url, storage_table_name, credential=credential\n",
    "    )\n",
    "    # output filename and file format\n",
    "    output_format = message[\"OutputFormat\"]\n",
    "    output_filename_timestamp = query_results_df[\"TimeGenerated\"].iloc[0]\n",
    "    output_filename_base = generate_output_filename_base(\n",
    "        message, output_filename_timestamp\n",
    "    )\n",
    "    full_output_filename, output_data = output_filename_and_format(\n",
    "        query_results_df, output_format, output_filename_base\n",
    "    )\n",
    "    # upload to blob storage\n",
    "    file_size = upload_file_to_storage(\n",
    "        container_client, full_output_filename, output_data\n",
    "    )\n",
    "    status = \"Success\"\n",
    "    # logging success to storage table\n",
    "    query_uuid = message[\"QueryUUID\"]\n",
    "    table_name = message[\"Table\"]\n",
    "    start_datetime = message[\"StartDatetime\"]\n",
    "    start_datetime = start_datetime.replace(\"T\", \" \").replace(\"Z\", \"\")\n",
    "    end_datetime = message[\"EndDatetime\"]\n",
    "    end_datetime = end_datetime.replace(\"T\", \" \").replace(\"Z\", \"\")\n",
    "    row_count = message[\"Count\"]\n",
    "    # generate unique row key\n",
    "    row_key = f\"{query_uuid}__{status}__{table_name}__\"\n",
    "    row_key += f\"{start_datetime}__{end_datetime}__{row_count}__\"\n",
    "    row_key += f\"{full_output_filename}__{file_size}\"\n",
    "    unique_row_sha256_hash = hashlib.sha256(row_key.encode()).hexdigest()\n",
    "    # response and logging to storage table\n",
    "    runtime_seconds = round(time.time() - start_time, 1)\n",
    "    time_generated = pd.Timestamp.today(\"UTC\").strftime(\"%Y-%m-%d %H:%M:%S.%f\")\n",
    "    return_message = {\n",
    "        \"PartitionKey\": query_uuid,\n",
    "        \"RowKey\": unique_row_sha256_hash,\n",
    "        \"Status\": status,\n",
    "        \"Table\": table_name,\n",
    "        \"StartDatetime\": start_datetime,\n",
    "        \"EndDatetime\": end_datetime,\n",
    "        \"RowCount\": row_count,\n",
    "        \"Filename\": full_output_filename,\n",
    "        \"FileSizeBytes\": file_size,\n",
    "        \"RuntimeSeconds\": runtime_seconds,\n",
    "        \"TimeGenerated\": time_generated,\n",
    "    }\n",
    "    table_client.upsert_entity(return_message, mode=UpdateMode.REPLACE)\n",
    "\n",
    "\n",
    "def process_queue_messages_loop(\n",
    "    credential: DefaultAzureCredential,\n",
    "    storage_queue_url: str,\n",
    "    storage_queue_name: str,\n",
    "    message_visibility_timeout_seconds: int = 600,\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    Processes Log Analytics query jobs/messages from a storage queue and exports to Blob Storage\n",
    "        note: credential requires Log Analytics Contributor, Storage Queue Data Contributor, and Storage Blob Data Contributor roles\n",
    "        note: takes ~150 seconds for a query with 500k rows and 10 columns to csv (100 seconds for parquet)\n",
    "        note: intended to be run interactively, for example, in a notebook or terminal\n",
    "        note: for production environment, use an azure function app\n",
    "    Args:\n",
    "        credential: azure default credential object\n",
    "        storage_queue_url: storage account queue url\n",
    "            format: \"https://{storage_account_name}.queue.core.windows.net/\"\n",
    "        storage_queue_name: name of queue\n",
    "        message_visibility_timeout_seconds: number of seconds for queue message visibility\n",
    "    Returns:\n",
    "        dict of results summary\n",
    "    \"\"\"\n",
    "    logging.info(\"Processing Queue Messages, press CTRL+C or interupt kernel to stop...\")\n",
    "    start_time = time.time()\n",
    "    # log analytics connection\n",
    "    # note: need to add Log Analytics Contributor role\n",
    "    log_client = LogsQueryClient(credential)\n",
    "    # storage queue connection\n",
    "    # note: need to add Storage Queue Data Contributor role\n",
    "    storage_queue_url_and_name = storage_queue_url + storage_queue_name\n",
    "    queue_client = QueueClient.from_queue_url(storage_queue_url_and_name, credential)\n",
    "    # process messages from queue until empty\n",
    "    successful_messages = 0\n",
    "    failed_messages = 0\n",
    "    try:\n",
    "        # loop through all messages in queue\n",
    "        while True:\n",
    "            # queue status\n",
    "            get_queue_properties = queue_client.get_queue_properties()\n",
    "            logging.info(\"Queue Status: %s\", get_queue_properties)\n",
    "            # get message\n",
    "            each_start_time = time.time()\n",
    "            queue_message = get_message_from_queue(\n",
    "                queue_client, message_visibility_timeout_seconds\n",
    "            )\n",
    "            if queue_message:\n",
    "                try:\n",
    "                    # extract content\n",
    "                    message_content = json.loads(queue_message.content)\n",
    "                    # process message: validate, query log analytics, upload to storage\n",
    "                    process_queue_message(log_client, message_content)\n",
    "                    # remove message from queue if successful\n",
    "                    delete_message_from_queue(queue_client, queue_message)\n",
    "                    successful_messages += 1\n",
    "                    runtime_calculation = round(time.time() - each_start_time, 1)\n",
    "                    logging.info(\"Runtime: %s\", runtime_calculation)\n",
    "                except Exception as e:\n",
    "                    logging.error(\n",
    "                        \"Unable to process message: %s %s\", queue_message.content, e\n",
    "                    )\n",
    "                    failed_messages += 1\n",
    "                    continue\n",
    "            # queue empty\n",
    "            else:\n",
    "                logging.info(\n",
    "                    \"Waiting for message visibility timeout (%s seconds)...\",\n",
    "                    message_visibility_timeout_seconds,\n",
    "                )\n",
    "                time.sleep(message_visibility_timeout_seconds + 60)\n",
    "                # check if queue still empty\n",
    "                if check_if_queue_empty_peek_message(queue_client):\n",
    "                    logging.info(\"No messages in queue\")\n",
    "                    break\n",
    "    # stop processing by keyboard interrupt\n",
    "    except KeyboardInterrupt:\n",
    "        logging.warning(\"Run was canceled manually by user\")\n",
    "    # return results\n",
    "    finally:\n",
    "        get_queue_properties = queue_client.get_queue_properties()\n",
    "        logging.info(\"Queue Status: %s\", get_queue_properties)\n",
    "        logging.info(\"Processing queue messages complete\")\n",
    "        return_result = {\n",
    "            \"successful_messages\": successful_messages,\n",
    "            \"failed_messages\": failed_messages,\n",
    "            \"runtime_seconds\": round(time.time() - start_time, 1),\n",
    "        }\n",
    "        return return_result  # pylint: disable=return-in-finally disable=lost-exception\n",
    "\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# storage blob\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "def upload_file_to_storage(\n",
    "    container_client: ContainerClient,\n",
    "    filename: str,\n",
    "    data: bytes | str,\n",
    "    azure_storage_connection_timeout_fix_seconds: int = 600,\n",
    ") -> int:\n",
    "    # note: need to use undocumented param connection_timeout to avoid timeout errors\n",
    "    # ref: https://stackoverflow.com/questions/65092741/solve-timeout-errors-on-file-uploads-with-new-azure-storage-blob-package\n",
    "    try:\n",
    "        blob_client = container_client.get_blob_client(filename)\n",
    "        blob_client.upload_blob(\n",
    "            data=data,\n",
    "            connection_timeout=azure_storage_connection_timeout_fix_seconds,\n",
    "            overwrite=True,\n",
    "        )\n",
    "        storage_account_name = container_client.account_name\n",
    "        container_name = container_client.container_name\n",
    "        logging.info(\n",
    "            \"Successfully Uploaded %s:%s/%s\",\n",
    "            storage_account_name,\n",
    "            container_name,\n",
    "            filename,\n",
    "        )\n",
    "        # file size\n",
    "        uploaded_file_metadata = list(container_client.list_blobs(filename))[0]\n",
    "        uploaded_file_size = uploaded_file_metadata.size\n",
    "        file_size_calculation_mb = uploaded_file_size / 1_000_000\n",
    "        logging.info(\"File Size: %s MBs\", file_size_calculation_mb)\n",
    "        return uploaded_file_size\n",
    "    except Exception as e:\n",
    "        logging.error(\"Unable to upload, %s, %s\", filename, e)\n",
    "        raise Exception(f\"Unable to upload, {filename}\") from e\n",
    "\n",
    "\n",
    "def download_blob(\n",
    "    filename: str,\n",
    "    credential: DefaultAzureCredential,\n",
    "    storage_blob_url: str,\n",
    "    storage_container_name: str,\n",
    ") -> pd.DataFrame:\n",
    "    # storage blob connection\n",
    "    # note: need to add Storage Blob Data Contributor role\n",
    "    container_client = ContainerClient(\n",
    "        storage_blob_url, storage_container_name, credential\n",
    "    )\n",
    "    # download data\n",
    "    blob_client = container_client.get_blob_client(filename)\n",
    "    downloaded_blob = blob_client.download_blob()\n",
    "    if filename.endswith(\".json\"):\n",
    "        stream = StringIO(downloaded_blob.content_as_text())\n",
    "        output_df = pd.read_json(stream, lines=True)\n",
    "    elif filename.endswith(\".csv\"):\n",
    "        stream = StringIO(downloaded_blob.content_as_text())\n",
    "        output_df = pd.read_csv(stream)\n",
    "    elif filename.endswith(\".parquet\"):\n",
    "        stream = BytesIO()\n",
    "        downloaded_blob.readinto(stream)\n",
    "        output_df = pd.read_parquet(stream, engine=\"pyarrow\")\n",
    "    else:\n",
    "        raise Exception(\"file extension not supported\")\n",
    "    return output_df\n",
    "\n",
    "\n",
    "def list_blobs_df(\n",
    "    credential: DefaultAzureCredential,\n",
    "    storage_blob_url: str,\n",
    "    storage_container_name: str,\n",
    ") -> pd.DataFrame:\n",
    "    # storage blob connection\n",
    "    # note: need to add Storage Blob Data Contributor role\n",
    "    container_client = ContainerClient(\n",
    "        storage_blob_url, storage_container_name, credential\n",
    "    )\n",
    "    # get blobs\n",
    "    results = []\n",
    "    for each_file in container_client.list_blobs():\n",
    "        each_name = each_file.name\n",
    "        each_size_mb = each_file.size / 1_000_000\n",
    "        each_date = each_file.creation_time\n",
    "        results.append([each_name, each_size_mb, each_date])\n",
    "    # convert to dataframe\n",
    "    df = pd.DataFrame(results, columns=[\"filename\", \"file_size_mb\", \"creation_time\"])\n",
    "    df = df.sort_values(\"creation_time\", ascending=False)\n",
    "    return df\n",
    "\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# storage table\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "def get_and_process_table_results(\n",
    "    credential: DefaultAzureCredential,\n",
    "    storage_table_url: str,\n",
    "    storage_table_query_name: str,\n",
    "    storage_table_process_name: str,\n",
    "    query_uuid: str,\n",
    ") -> dict[str, pd.DataFrame]:\n",
    "    cols_to_rename = {\"PartitionKey\": \"QueryUUID\"}\n",
    "    return_dfs = {\n",
    "        \"query_results_df\": pd.DataFrame(),\n",
    "        \"process_results_df\": pd.DataFrame(),\n",
    "        \"success_process_results_df\": pd.DataFrame(),\n",
    "        \"failed_process_results_df\": pd.DataFrame(),\n",
    "    }\n",
    "    # table connections\n",
    "    table_client_query = TableClient(\n",
    "        storage_table_url, storage_table_query_name, credential=credential\n",
    "    )\n",
    "    table_client_process = TableClient(\n",
    "        storage_table_url, storage_table_process_name, credential=credential\n",
    "    )\n",
    "    # get results from azure storage tables\n",
    "    search_odata_string = f\"PartitionKey eq '{query_uuid}'\"\n",
    "    query_results = table_client_query.query_entities(search_odata_string)\n",
    "    process_results = table_client_process.query_entities(search_odata_string)\n",
    "    # query results\n",
    "    query_results_df_raw = pd.DataFrame(query_results)\n",
    "    if query_results_df_raw.shape[0] == 0:\n",
    "        logging.info(\"Query UUID %s not found in query logs\", query_uuid)\n",
    "        return return_dfs\n",
    "    else:\n",
    "        query_results_df = query_results_df_raw.rename(columns=cols_to_rename)\n",
    "        return_dfs[\"query_results_df\"] = query_results_df\n",
    "    # process results\n",
    "    process_results_df_raw = pd.DataFrame(process_results)\n",
    "    if process_results_df_raw.shape[0] == 0:\n",
    "        logging.info(\"Query UUID %s not found in process logs\", query_uuid)\n",
    "        return return_dfs\n",
    "    else:\n",
    "        process_results_df = process_results_df_raw.rename(columns=cols_to_rename)\n",
    "        return_dfs[\"process_results_df\"] = process_results_df\n",
    "        # split\n",
    "        success_mask = process_results_df.Status == \"Success\"\n",
    "        failed_mask = process_results_df.Status == \"Failed\"\n",
    "        success_process_results_df = process_results_df.loc[success_mask]\n",
    "        failed_process_results_df = process_results_df.loc[failed_mask]\n",
    "        return_dfs[\"success_process_results_df\"] = success_process_results_df\n",
    "        return_dfs[\"failed_process_results_df\"] = failed_process_results_df\n",
    "        return return_dfs\n",
    "\n",
    "\n",
    "def calculate_runtime_since_query_submit(\n",
    "    query_results_df: pd.DataFrame, process_results_df: pd.DataFrame\n",
    ") -> int:\n",
    "    # change column types to datetime\n",
    "    query_results_df_copy = query_results_df.copy()\n",
    "    query_results_df_copy[\"TimeGenerated\"] = pd.to_datetime(\n",
    "        query_results_df.TimeGenerated\n",
    "    )\n",
    "    process_results_df_copy = process_results_df.copy()\n",
    "    process_results_df_copy[\"TimeGenerated\"] = pd.to_datetime(\n",
    "        process_results_df_copy.TimeGenerated\n",
    "    )\n",
    "    # calculate difference between first submit and last processed log\n",
    "    query_submit_datetime = query_results_df_copy[\"TimeGenerated\"].min()\n",
    "    last_processing_datetime = process_results_df_copy[\"TimeGenerated\"].max()\n",
    "    time_since_query = last_processing_datetime - query_submit_datetime\n",
    "    time_since_query_seconds = time_since_query.total_seconds()\n",
    "    return time_since_query_seconds\n",
    "\n",
    "\n",
    "def calculate_processing_status_and_percent(\n",
    "    number_of_successful_subqueries: int,\n",
    "    number_of_subqueries: int,\n",
    "    total_success_row_count: int,\n",
    "    query_total_row_count: int,\n",
    ") -> tuple[str, float]:\n",
    "    if (\n",
    "        number_of_successful_subqueries == number_of_subqueries\n",
    "        and total_success_row_count == query_total_row_count\n",
    "    ):\n",
    "        processing_status = \"Complete\"\n",
    "    else:\n",
    "        processing_status = \"Partial\"\n",
    "    percent_complete = (number_of_successful_subqueries / number_of_subqueries) * 100\n",
    "    percent_complete = round(percent_complete, 1)\n",
    "    return processing_status, percent_complete\n",
    "\n",
    "\n",
    "def calculate_file_size(\n",
    "    filesize_units: str,\n",
    "    total_success_bytes: int,\n",
    ") -> tuple[float, str]:\n",
    "    if filesize_units == \"GB\":\n",
    "        divisor = 1_000_000_000\n",
    "        success_total_size = float(round(total_success_bytes / divisor, 3))\n",
    "        file_units = \"GB\"\n",
    "    elif filesize_units == \"TB\":\n",
    "        divisor = 1_000_000_000_000\n",
    "        success_total_size = float(round(total_success_bytes / divisor, 3))\n",
    "        file_units = \"TB\"\n",
    "    # defaults to MB\n",
    "    else:\n",
    "        divisor = 1_000_000\n",
    "        success_total_size = float(round(total_success_bytes / divisor, 3))\n",
    "        file_units = \"MB\"\n",
    "    return success_total_size, file_units\n",
    "\n",
    "\n",
    "def calculate_time_remaining_estimate(\n",
    "    processing_status: str,\n",
    "    number_of_successful_subqueries: int,\n",
    "    percent_complete: float,\n",
    "    time_since_query_seconds: int,\n",
    ") -> float | None:\n",
    "    if processing_status == \"Complete\":\n",
    "        time_remaining_seconds = 0.0\n",
    "    else:\n",
    "        if number_of_successful_subqueries > 0:\n",
    "            percent_remaining = 100 - percent_complete\n",
    "            time_remaining_seconds = (\n",
    "                time_since_query_seconds * percent_remaining / percent_complete\n",
    "            )\n",
    "            time_remaining_seconds = float(round(time_remaining_seconds, 1))\n",
    "        else:\n",
    "            time_remaining_seconds = None\n",
    "    return time_remaining_seconds\n",
    "\n",
    "\n",
    "def get_status(\n",
    "    credential: DefaultAzureCredential,\n",
    "    query_uuid: str,\n",
    "    storage_table_url: str,\n",
    "    storage_table_query_name: str,\n",
    "    storage_table_process_name: str,\n",
    "    return_failures: bool = True,\n",
    "    filesize_units: str = \"GB\",\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    Gets status of submitted query\n",
    "    Args:\n",
    "        query_uuid: query uuid or \"PartitionKey\"\n",
    "            format: \"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX\"\n",
    "        storage_table_url: storage table url\n",
    "            format: \"https://{storage_account_name}.table.core.windows.net/\"\n",
    "        storage_table_query_name: name of storage table for query logs\n",
    "        storage_table_process_name: name of storage table for process logs\n",
    "        return_failures: will return details on failed jobs/messages\n",
    "        filesize_units: \"MB\", \"GB\", or \"TB\"\n",
    "    Returns:\n",
    "        dict with high-level status properties\n",
    "    \"\"\"\n",
    "    # get table logs dataframes\n",
    "    table_dfs = get_and_process_table_results(\n",
    "        credential,\n",
    "        storage_table_url,\n",
    "        storage_table_query_name,\n",
    "        storage_table_process_name,\n",
    "        query_uuid,\n",
    "    )\n",
    "    query_results_df = table_dfs[\"query_results_df\"]\n",
    "    process_results_df = table_dfs[\"process_results_df\"]\n",
    "    success_process_results_df = table_dfs[\"success_process_results_df\"]\n",
    "    failed_process_results_df = table_dfs[\"failed_process_results_df\"]\n",
    "    # query results\n",
    "    results = {\"query_uuid\": query_uuid}\n",
    "    no_results_status = \"No Records Found\"\n",
    "    if query_results_df.empty:\n",
    "        results[\"query_submit_status\"] = no_results_status\n",
    "        results[\"query_processing_status\"] = no_results_status\n",
    "        return results\n",
    "    if len(query_results_df.Status.unique()) == 1:\n",
    "        query_results_status = str(query_results_df.Status.unique()[0])\n",
    "    else:\n",
    "        query_results_status = str(query_results_df.Status.value_counts().to_dict())\n",
    "        query_results_status = query_results_status.replace(\"{\", \"\").replace(\"}\", \"\")\n",
    "        query_results_status = query_results_status.replace(\"'\", \"\")\n",
    "    query_submit_status = query_results_status\n",
    "    number_of_query_splits = query_results_df.shape[0]\n",
    "    number_of_subqueries = query_results_df.MessagesSentToQueue.sum()\n",
    "    query_total_row_count = query_results_df.TotalRowCount.sum()\n",
    "    results[\"query_submit_status\"] = query_submit_status\n",
    "    results[\"query_submit_splits\"] = number_of_query_splits\n",
    "    results[\"number_of_subqueries\"] = int(number_of_subqueries)\n",
    "    results[\"query_total_row_count\"] = int(query_total_row_count)\n",
    "    # process results\n",
    "    if process_results_df.empty:\n",
    "        results[\"query_processing_status\"] = no_results_status\n",
    "        return results\n",
    "    number_of_successful_subqueries = success_process_results_df.shape[0]\n",
    "    number_of_failed_subqueries = failed_process_results_df.shape[0]\n",
    "    total_success_bytes = success_process_results_df.FileSizeBytes.sum()\n",
    "    total_success_row_count = success_process_results_df.RowCount.sum()\n",
    "    time_since_query_seconds = calculate_runtime_since_query_submit(\n",
    "        query_results_df, process_results_df\n",
    "    )\n",
    "    processing_status, percent_complete = calculate_processing_status_and_percent(\n",
    "        number_of_successful_subqueries,\n",
    "        number_of_subqueries,\n",
    "        total_success_row_count,\n",
    "        query_total_row_count,\n",
    "    )\n",
    "    success_total_size, file_units = calculate_file_size(\n",
    "        filesize_units, total_success_bytes\n",
    "    )\n",
    "    time_remaining_seconds = calculate_time_remaining_estimate(\n",
    "        processing_status,\n",
    "        number_of_successful_subqueries,\n",
    "        percent_complete,\n",
    "        time_since_query_seconds,\n",
    "    )\n",
    "    results[\"query_processing_status\"] = processing_status\n",
    "    results[\"processing_percent_complete\"] = float(percent_complete)\n",
    "    results[\"number_of_subqueries_success\"] = number_of_successful_subqueries\n",
    "    results[\"number_of_subqueries_failed\"] = number_of_failed_subqueries\n",
    "    results[\"output_total_row_count\"] = int(total_success_row_count)\n",
    "    results[\"output_file_size\"] = success_total_size\n",
    "    results[\"output_file_units\"] = file_units\n",
    "    results[\"runtime_since_submit_seconds\"] = round(time_since_query_seconds, 1)\n",
    "    results[\"processing_estimated_time_remaining_seconds\"] = time_remaining_seconds\n",
    "    # failures\n",
    "    if return_failures and failed_process_results_df.shape[0] > 0:\n",
    "        export_cols = [\n",
    "            \"Table\",\n",
    "            \"StartDatetime\",\n",
    "            \"EndDatetime\",\n",
    "            \"RowCount\",\n",
    "        ]\n",
    "        rename_columns = {\n",
    "            \"Table\": \"table\",\n",
    "            \"StartDatetime\": \"start_datetime\",\n",
    "            \"EndDatetime\": \"end_datetime\",\n",
    "            \"RowCount\": \"row_count\",\n",
    "        }\n",
    "        export_df = failed_process_results_df[export_cols]\n",
    "        export_df_renamed = export_df.rename(columns=rename_columns)\n",
    "        results[\"failures\"] = export_df_renamed.to_dict(orient=\"records\")\n",
    "    else:\n",
    "        results[\"failures\"] = [{}]\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd4ae643-12f8-4323-bfae-a9fc8cb81413",
   "metadata": {},
   "source": [
    "## Authentication "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc3474c1-8bff-4377-91b9-087f331ef1a7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Option 1. service principal\n",
    "# os.environ[\"AZURE_CLIENT_ID\"] = \"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX\"\n",
    "# os.environ[\"AZURE_TENANT_ID\"] = \"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX\"\n",
    "# os.environ[\"AZURE_CLIENT_SECRET\"] = \"XXXXXXXXXXXXXXXXXXXXXXXXXXXX\"\n",
    "# Option 2. command line -> AAD/Entra account\n",
    "!az login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a367f3e2-12ff-4f1d-a244-3677869a0853",
   "metadata": {},
   "outputs": [],
   "source": [
    "credential = DefaultAzureCredential()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d3b2ee9-296b-401b-bce3-e7c340d6e580",
   "metadata": {},
   "source": [
    "## 1. Ingest Test Data into Log Analytics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15ccb098-f086-45dc-8b64-4c92fc9017ad",
   "metadata": {},
   "source": [
    "### 1.1 Inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "398b493c-af04-4d03-9f35-e51ac26a9d48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# connection\n",
    "log_analytics_data_collection_endpoint = \"https://XXXXXXXXX.XXXX.ingest.monitor.azure.com\"\n",
    "log_analytics_data_collection_rule_id = \"dcr-XXXXXXXXXXXXXXXXXXXXXXXXXX\"\n",
    "log_analytics_data_collection_stream_name = \"Custom-XXXXXXXXXXXXXXXXXXX_CL\"\n",
    "storage_table_url = \"https://XXXXXXXXXXXXXXXXXXXXXXXX.table.core.windows.net/\"\n",
    "storage_table_ingest_name = \"XXXXXXXXXXXX\"\n",
    "# params\n",
    "start_datetime = \"2024-04-04 00:00:00.000000\"\n",
    "timedelta_seconds = 0.000_3\n",
    "number_of_rows = 100_000_000"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5272e9cb-d131-4f8b-a3fc-5cf1743487cb",
   "metadata": {},
   "source": [
    "### 1.2 Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e3a3a88-b799-4df1-8928-c91dc5d00721",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ingest_results = generate_and_ingest_test_data(\n",
    "    credential,\n",
    "    log_analytics_data_collection_endpoint,\n",
    "    log_analytics_data_collection_rule_id,\n",
    "    log_analytics_data_collection_stream_name,\n",
    "    storage_table_url,\n",
    "    storage_table_ingest_name,\n",
    "    start_datetime,\n",
    "    timedelta_seconds,\n",
    "    number_of_rows,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bf15469-7375-451f-bb22-e2aaa17a89d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "ingest_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52f17185-b1b7-4ffd-8c50-e5b5ec3edb72",
   "metadata": {},
   "source": [
    "### 1.3 Verify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0300bb7d-7cf7-46c7-8f56-dfb4548c745b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# storage table log\n",
    "storage_table_url = \"https://XXXXXXXXXXXXXXXXX.table.core.windows.net/\"\n",
    "storage_table_ingest_name = \"XXXXXXXXXXXXXXX\"\n",
    "storage_table_url_and_name = storage_table_url + storage_table_ingest_name\n",
    "ingest_log = TableClient.from_table_url(storage_table_url_and_name, credential=credential)\n",
    "ingest_log_df = pd.DataFrame(ingest_log.list_entities())\n",
    "ingest_log_df[\n",
    "    [\n",
    "        \"TimeGenerated\",\n",
    "        \"PartitionKey\",\n",
    "        \"Status\",\n",
    "        \"StartDatetime\",\n",
    "        \"EndDatetime\",\n",
    "        \"NumberColumns\",\n",
    "        \"RowsGenerated\",\n",
    "        \"RowsIngested\",\n",
    "        \"ValidDatetimeRange\",\n",
    "        \"RuntimeSeconds\",\n",
    "    ]\n",
    "].sort_values(\"TimeGenerated\", ascending=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07f215b1-ba42-4a5f-a3cb-3844eed28208",
   "metadata": {},
   "outputs": [],
   "source": [
    "# query log analytics to confirm ingest\n",
    "log_analytics_workspace_id = \"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX\"\n",
    "log_analytics_table_name = log_analytics_data_collection_stream_name[7:]\n",
    "test_kql_query = f\"\"\"\n",
    "let TABLE_NAME = \"{log_analytics_table_name}\";\n",
    "table(TABLE_NAME)\n",
    "| project-away TenantId, Type, _ResourceId\n",
    "| sort by TimeGenerated desc\n",
    "| take 100\n",
    "\"\"\"\n",
    "query_log_analytics_connection_request(\n",
    "    credential,\n",
    "    log_analytics_workspace_id,\n",
    "    test_kql_query,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f0411e7-a286-4b3a-b9f1-7bba4ed94ed1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# query log analytics to check count\n",
    "verify_start_datetime = \"2024-04-04 00:00:00\"\n",
    "verify_end_datetime = \"2024-04-05 00:00:00\"\n",
    "log_analytics_table_name = log_analytics_data_collection_stream_name[7:]\n",
    "test_kql_query = f\"\"\"\n",
    "let TABLE_NAME = \"{log_analytics_table_name}\";\n",
    "let START_DATETIME = datetime({verify_start_datetime});\n",
    "let END_DATETIME = datetime({verify_end_datetime});\n",
    "table(TABLE_NAME)\n",
    "| project TimeGenerated\n",
    "| where (TimeGenerated >= START_DATETIME) and (TimeGenerated < END_DATETIME)\n",
    "| count\n",
    "\"\"\"\n",
    "query_log_analytics_connection_request(\n",
    "    credential,\n",
    "    log_analytics_workspace_id,\n",
    "    test_kql_query,\n",
    ").Count[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d71ef7f-1efe-4909-9ced-252d62bca802",
   "metadata": {},
   "source": [
    "## 2. Split Query and Send to Queue"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61a67a6a-5a57-4f1a-952f-86a45df6afd5",
   "metadata": {},
   "source": [
    "### 2.1 Inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "930f8c3f-f9a8-454b-a73b-6de1e467d4a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# connections\n",
    "subscription_id = \"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX\"\n",
    "resource_group_name = \"XXXXXXXXXXXXXXXXXXXXX\"\n",
    "log_analytics_worksapce_name = \"XXXXXXXXXXXXXXXX\"\n",
    "log_analytics_workspace_id = \"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX\"\n",
    "storage_queue_url = \"https://XXXXXXXXXXXXXXXXXXXXXX.queue.core.windows.net/\"\n",
    "storage_queue_name = \"XXXXXXXXXXXXX\"\n",
    "storage_blob_url = \"https://XXXXXXXXXXXXXXXXXXX.blob.core.windows.net/\"\n",
    "storage_blob_container_name = \"XXXXXXXXXXXXX\"\n",
    "storage_table_url = \"https://XXXXXXXXXXXXXXXXXXXXX.table.core.windows.net/\"\n",
    "storage_table_query_name = \"XXXXXXXXXXX\"\n",
    "storage_table_process_name = \"XXXXXXXXXXXX\"\n",
    "# params\n",
    "table_names_and_columns = {\n",
    "    \"XXXXXXXXXXXXXXXX_CL\": [\n",
    "        \"TimeGenerated\",\n",
    "        \"DataColumn1\",\n",
    "        \"DataColumn2\",\n",
    "        \"DataColumn3\",\n",
    "        \"DataColumn4\",\n",
    "        \"DataColumn5\",\n",
    "        \"DataColumn6\",\n",
    "        \"DataColumn7\",\n",
    "        \"DataColumn8\",\n",
    "        \"DataColumn9\",\n",
    "    ]\n",
    "}\n",
    "start_datetime = \"2024-04-04 00:00:00\"\n",
    "end_datetime = \"2024-04-05 00:00:00\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48de5a23-c068-4bdb-b88a-090d63729575",
   "metadata": {},
   "source": [
    "### 2.2 Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c014a2f-5daf-4dc4-a6cf-ef1430768323",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_uuid = str(uuid.uuid4())\n",
    "print(f\"Query UUID: {query_uuid}\")\n",
    "query_results = query_log_analytics_send_to_queue(\n",
    "    query_uuid,\n",
    "    credential,\n",
    "    subscription_id,\n",
    "    resource_group_name,\n",
    "    log_analytics_worksapce_name,\n",
    "    log_analytics_workspace_id,\n",
    "    storage_queue_url,\n",
    "    storage_queue_name,\n",
    "    storage_blob_url,\n",
    "    storage_blob_container_name,\n",
    "    storage_table_url,\n",
    "    storage_table_query_name,\n",
    "    storage_table_process_name,\n",
    "    table_names_and_columns,\n",
    "    start_datetime,\n",
    "    end_datetime,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a0b2280-0e74-4454-8782-ce5e73be9c80",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b23730de-8f9a-422f-9494-51019879dbd4",
   "metadata": {},
   "source": [
    "### 2.3 Verify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1db7f67-546e-4cce-b036-e1043a42bff6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# storage table log\n",
    "storage_table_url = \"https://XXXXXXXXXXXXXXXX.table.core.windows.net/\"\n",
    "storage_table_query_name = \"XXXXXXXXXXX\"\n",
    "storage_table_url_and_name = storage_table_url + storage_table_query_name\n",
    "query_log = TableClient.from_table_url(storage_table_url_and_name, credential=credential)\n",
    "query_log_df = pd.DataFrame(query_log.list_entities())\n",
    "query_log_df[\n",
    "    [\n",
    "        \"TimeGenerated\",\n",
    "        \"PartitionKey\",\n",
    "        \"Status\",\n",
    "        \"Tables\",\n",
    "        \"StartDatetime\",\n",
    "        \"EndDatetime\",\n",
    "        \"TotalRowCount\",\n",
    "        \"MessagesGenerated\",\n",
    "        \"MessagesSentToQueue\",\n",
    "        \"RuntimeSeconds\",\n",
    "    ]\n",
    "].sort_values(\"TimeGenerated\", ascending=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b341e9fd-f695-47e6-8de7-0886b0c51c97",
   "metadata": {},
   "outputs": [],
   "source": [
    "# queue size\n",
    "storage_queue_url = \"https://XXXXXXXXXXXXXXXXXXX.queue.core.windows.net/\"\n",
    "storage_queue_name = \"XXXXXXXXXXXXXXXXXX\"\n",
    "storage_queue_url_and_name = storage_queue_url + storage_queue_name\n",
    "QueueClient.from_queue_url(storage_queue_url_and_name, credential).get_queue_properties()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be11e3da-192c-4db3-9749-6b49df42f9ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# preview queue messages\n",
    "storage_queue_url_and_name = storage_queue_url + storage_queue_name\n",
    "QueueClient.from_queue_url(storage_queue_url_and_name, credential).peek_messages(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e9d0dec-584b-4f2a-970f-90c9dfef76df",
   "metadata": {},
   "source": [
    "## 3. Process Queue"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c309d590-7e38-43ec-b258-61d144dca816",
   "metadata": {},
   "source": [
    "### 3.1 Inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c005e22f-2107-4be1-96bb-5d0d611da181",
   "metadata": {},
   "outputs": [],
   "source": [
    "# connections\n",
    "storage_queue_url = \"https://XXXXXXXXXXXXXXXXXXXX.queue.core.windows.net/\"\n",
    "storage_queue_name = \"XXXXXXXXXXXX\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c0aa1fb-cf64-4fe0-b653-c1fe2a9b6336",
   "metadata": {},
   "source": [
    "### 3.2 Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe640a60-d274-44df-b8fa-b6069274806d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "process_results = process_queue_messages_loop(\n",
    "    credential,\n",
    "    storage_queue_url,\n",
    "    storage_queue_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63f84b53-8365-4ae8-89b5-40b44a5aa78b",
   "metadata": {},
   "outputs": [],
   "source": [
    "process_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ae30247-1a7e-4adb-9fe9-527b2170cb1e",
   "metadata": {},
   "source": [
    "### 3.3 Verify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4274fc46-6d88-4b93-911e-2395770fcd34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# storage table log\n",
    "storage_table_url = \"https://XXXXXXXXXXXXXXXXXXX.table.core.windows.net/\"\n",
    "storage_table_process_name = \"XXXXXXXXXXXXX\"\n",
    "storage_table_url_and_name = storage_table_url + storage_table_process_name\n",
    "query_log = TableClient.from_table_url(storage_table_url_and_name, credential=credential)\n",
    "query_log_df = pd.DataFrame(query_log.list_entities())\n",
    "query_log_df[\n",
    "    [\n",
    "        \"TimeGenerated\",\n",
    "        \"PartitionKey\",\n",
    "        \"SubQuery\",\n",
    "        \"Status\",\n",
    "        \"Table\",\n",
    "        \"StartDatetime\",\n",
    "        \"EndDatetime\",\n",
    "        \"RowCount\",\n",
    "        \"Filename\",\n",
    "        \"RuntimeSeconds\",\n",
    "    ]\n",
    "].sort_values(\"TimeGenerated\", ascending=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d482edd5-cb5c-4cfe-8d9c-7f5527961a66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# list files\n",
    "storage_blob_url = \"https://XXXXXXXXXXXXXXXXXXXXXX.blob.core.windows.net/\"\n",
    "storage_blob_container_name = \"XXXXXXXXXXXXXXX\"\n",
    "list_blobs_df(credential, storage_blob_url, storage_blob_container_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1362b0ae-1ccc-4424-a135-4779ef786109",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read file\n",
    "file_list = list_blobs_df(credential, storage_blob_url, storage_blob_container_name)\n",
    "most_recent_filename = file_list.sort_values(\"creation_time\")[\"filename\"].iloc[0]\n",
    "download_blob(\n",
    "    most_recent_filename, credential, storage_blob_url, storage_blob_container_name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5008557-aa47-40cb-8f4f-6afcb23eac70",
   "metadata": {},
   "source": [
    "### 3.4 Get Status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b2def41-3a19-4075-8b8c-2216786482f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# connection\n",
    "storage_table_url = \"https://XXXXXXXXXXXXXXXXXXX.table.core.windows.net/\"\n",
    "storage_table_query_name = \"XXXXXXXXXXX\"\n",
    "storage_table_process_name = \"XXXXXXXXXX\"\n",
    "# params\n",
    "query_uuid = \"XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a9f0039-fb27-467a-b13a-00fde51cb74f",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_status(\n",
    "    credential,\n",
    "    query_uuid,\n",
    "    storage_table_url,\n",
    "    storage_table_query_name,\n",
    "    storage_table_process_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24eaa4cc-4aa3-4f57-9f82-8c9e7eb0c66d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}