Separate content of csm-data and coal for adx_send_data

Author: neomatamune

cosmotech/coal/azure/adx/__init__.py

@@ -5,16 +5,19 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
-from cosmotech.coal.azure.adx.auth import create_kusto_client, create_ingest_client
+from cosmotech.coal.azure.adx.auth import create_kusto_client, create_ingest_client, initialize_clients
 from cosmotech.coal.azure.adx.query import run_query, run_command_query
 from cosmotech.coal.azure.adx.ingestion import (
     ingest_dataframe,
     send_to_adx,
     check_ingestion_status,
+    monitor_ingestion,
+    handle_failures,
     IngestionStatus,
 )
-from cosmotech.coal.azure.adx.tables import table_exists, create_table
-from cosmotech.coal.azure.adx.utils import type_mapping
+from cosmotech.coal.azure.adx.tables import table_exists, create_table, check_and_create_table, _drop_by_tag
+from cosmotech.coal.azure.adx.utils import type_mapping, create_column_mapping
+from cosmotech.coal.azure.adx.store import send_pyarrow_table_to_adx, send_table_data, process_tables, send_store_to_adx
 from cosmotech.coal.azure.adx.runner import (
     prepare_csv_content,
     construct_create_query,

cosmotech/coal/azure/adx/auth.py

@@ -6,7 +6,7 @@
 # specifically authorized by written means by Cosmo Tech.
 
 import os
-from typing import Union, Optional
+from typing import Union, Optional, Tuple
 
 from azure.kusto.data import KustoClient, KustoConnectionStringBuilder
 from azure.kusto.ingest import QueuedIngestClient
@@ -87,7 +87,24 @@ def create_ingest_client(
     return QueuedIngestClient(kcsb)
 
 
-def get_cluster_urls(cluster_name: str, cluster_region: str) -> tuple[str, str]:
+def initialize_clients(adx_uri: str, adx_ingest_uri: str) -> Tuple[KustoClient, QueuedIngestClient]:
+    """
+    Initialize and return the Kusto and ingest clients.
+
+    Args:
+        adx_uri: The Azure Data Explorer resource URI
+        adx_ingest_uri: The Azure Data Explorer resource ingest URI
+
+    Returns:
+        tuple: (kusto_client, ingest_client)
+    """
+    LOGGER.debug("Initializing clients")
+    kusto_client = create_kusto_client(adx_uri)
+    ingest_client = create_ingest_client(adx_ingest_uri)
+    return kusto_client, ingest_client
+
+
+def get_cluster_urls(cluster_name: str, cluster_region: str) -> Tuple[str, str]:
     """
     Generate cluster and ingest URLs from cluster name and region.
 

cosmotech/coal/azure/adx/ingestion.py

@@ -12,8 +12,10 @@
 from typing import Optional
 from typing import Tuple
 
+import os
 import pandas as pd
 import time
+import tqdm
 from azure.kusto.data import KustoClient
 from azure.kusto.data.data_format import DataFormat
 from azure.kusto.ingest import IngestionProperties
@@ -24,7 +26,7 @@
 from azure.kusto.ingest.status import SuccessMessage
 from cosmotech.orchestrator.utils.translate import T
 
-from cosmotech.coal.azure.adx.tables import create_table
+from cosmotech.coal.azure.adx.tables import create_table, _drop_by_tag
 from cosmotech.coal.azure.adx.utils import type_mapping
 from cosmotech.coal.utils.logger import LOGGER
 
@@ -221,6 +223,86 @@ def get_messages(queues):
         yield source_id, _ingest_status[source_id]
 
 
+def monitor_ingestion(
+    ingest_client: QueuedIngestClient, source_ids: List[str], table_ingestion_id_mapping: Dict[str, str]
+) -> bool:
+    """
+    Monitor the ingestion process with progress reporting.
+
+    Args:
+        ingest_client: The ingest client
+        source_ids: List of source IDs to monitor
+        table_ingestion_id_mapping: Mapping of source IDs to table names
+
+    Returns:
+        bool: True if any failures occurred, False otherwise
+    """
+    has_failures = False
+    source_ids_copy = source_ids.copy()
+
+    LOGGER.info("Waiting for ingestion of data to finish")
+
+    with tqdm.tqdm(desc="Ingestion status", total=len(source_ids_copy)) as pbar:
+        while any(
+            list(
+                map(
+                    lambda _status: _status[1] in (IngestionStatus.QUEUED, IngestionStatus.UNKNOWN),
+                    results := list(check_ingestion_status(ingest_client, source_ids_copy)),
+                )
+            )
+        ):
+            # Check for failures
+            for ingestion_id, ingestion_status in results:
+                if ingestion_status == IngestionStatus.FAILURE:
+                    LOGGER.error(
+                        f"Ingestion {ingestion_id} failed for table {table_ingestion_id_mapping.get(ingestion_id)}"
+                    )
+                    has_failures = True
+
+            cleared_ids = list(
+                result for result in results if result[1] not in (IngestionStatus.QUEUED, IngestionStatus.UNKNOWN)
+            )
+
+            for ingestion_id, ingestion_status in cleared_ids:
+                pbar.update(1)
+                source_ids_copy.remove(ingestion_id)
+
+            time.sleep(1)
+            if os.environ.get("CSM_USE_RICH", "False").lower() in ("true", "1", "yes", "t", "y"):
+                pbar.refresh()
+        else:
+            for ingestion_id, ingestion_status in results:
+                if ingestion_status == IngestionStatus.FAILURE:
+                    LOGGER.error(
+                        f"Ingestion {ingestion_id} failed for table {table_ingestion_id_mapping.get(ingestion_id)}"
+                    )
+                    has_failures = True
+        pbar.update(len(source_ids_copy))
+
+    LOGGER.info("All data ingestion attempts completed")
+    return has_failures
+
+
+def handle_failures(kusto_client: KustoClient, database: str, operation_tag: str, has_failures: bool) -> bool:
+    """
+    Handle any failures and perform rollbacks if needed.
+
+    Args:
+        kusto_client: The Kusto client
+        database: The database name
+        operation_tag: The operation tag for tracking
+        has_failures: Whether any failures occurred
+
+    Returns:
+        bool: True if the process should abort, False otherwise
+    """
+    if has_failures:
+        LOGGER.warning(f"Failures detected during ingestion - dropping data with tag: {operation_tag}")
+        _drop_by_tag(kusto_client, database, operation_tag)
+        return True
+    return False
+
+
 def clear_ingestion_status_queues(client: QueuedIngestClient, confirmation: bool = False):
     """
     Clear all data in the ingestion status queues.

cosmotech/coal/azure/adx/store.py

@@ -8,23 +8,87 @@
 import os
 import tempfile
 import uuid
-from typing import Optional
+from typing import Optional, List, Dict, Tuple, Union, Any
 
 import pyarrow
 import pyarrow.csv as pc
 import time
+from azure.kusto.data import KustoClient
 from azure.kusto.data.data_format import DataFormat
 from azure.kusto.ingest import IngestionProperties
 from azure.kusto.ingest import QueuedIngestClient
 from azure.kusto.ingest import ReportLevel
 from cosmotech.orchestrator.utils.translate import T
 from time import perf_counter
 
+from cosmotech.coal.azure.adx.tables import check_and_create_table, _drop_by_tag
+from cosmotech.coal.azure.adx.auth import initialize_clients
+from cosmotech.coal.azure.adx.ingestion import monitor_ingestion, handle_failures
 from cosmotech.coal.store.store import Store
 from cosmotech.coal.utils.logger import LOGGER
 from cosmotech.coal.utils.postgresql import send_pyarrow_table_to_postgresql
 
 
+def send_table_data(
+    ingest_client: QueuedIngestClient, database: str, table_name: str, data: pyarrow.Table, operation_tag: str
+) -> Tuple[str, str]:
+    """
+    Send a PyArrow table to ADX.
+
+    Args:
+        ingest_client: The ingest client
+        database: The database name
+        table_name: The table name
+        data: The PyArrow table data
+        operation_tag: The operation tag for tracking
+
+    Returns:
+        tuple: (source_id, table_name)
+    """
+    LOGGER.debug(f"Sending data to the table {table_name}")
+    result = send_pyarrow_table_to_adx(ingest_client, database, table_name, data, operation_tag)
+    return result.source_id, table_name
+
+
+def process_tables(
+    store: Store, kusto_client: KustoClient, ingest_client: QueuedIngestClient, database: str, operation_tag: str
+) -> Tuple[List[str], Dict[str, str]]:
+    """
+    Process all tables in the store.
+
+    Args:
+        store: The data store
+        kusto_client: The Kusto client
+        ingest_client: The ingest client
+        database: The database name
+        operation_tag: The operation tag for tracking
+
+    Returns:
+        tuple: (source_ids, table_ingestion_id_mapping)
+    """
+    source_ids = []
+    table_ingestion_id_mapping = dict()
+
+    LOGGER.debug("Listing tables")
+    table_list = list(store.list_tables())
+
+    for target_table_name in table_list:
+        LOGGER.info(f"Working on table: {target_table_name}")
+        data = store.get_table(target_table_name)
+
+        if data.num_rows < 1:
+            LOGGER.warning(f"Table {target_table_name} has no rows - skipping it")
+            continue
+
+        check_and_create_table(kusto_client, database, target_table_name, data)
+
+        source_id, _ = send_table_data(ingest_client, database, target_table_name, data, operation_tag)
+        source_ids.append(source_id)
+        table_ingestion_id_mapping[source_id] = target_table_name
+
+    return source_ids, table_ingestion_id_mapping
+
+
 def send_pyarrow_table_to_adx(
     client: QueuedIngestClient,
     database: str,
@@ -52,6 +116,68 @@ def send_pyarrow_table_to_adx(
         os.unlink(temp_file_path)
 
 
+def send_store_to_adx(
+    adx_uri: str,
+    adx_ingest_uri: str,
+    database_name: str,
+    wait: bool = False,
+    tag: Optional[str] = None,
+    store_location: Optional[str] = None,
+) -> Union[bool, Any]:
+    """
+    Send data from the store to Azure Data Explorer.
+
+    Args:
+        adx_uri: The Azure Data Explorer resource URI
+        adx_ingest_uri: The Azure Data Explorer resource ingest URI
+        database_name: The database name
+        wait: Whether to wait for ingestion to complete
+        tag: The operation tag for tracking (will generate a unique one if not provided)
+        store_location: Optional store location (uses default if not provided)
+
+    Returns:
+        bool: True if successful, False otherwise
+    """
+    # Generate a unique operation tag if none provided
+    operation_tag = tag or f"op-{str(uuid.uuid4())}"
+    LOGGER.debug(f"Starting ingestion operation with tag: {operation_tag}")
+
+    # Initialize clients
+    kusto_client, ingest_client = initialize_clients(adx_uri, adx_ingest_uri)
+    database = database_name
+
+    # Load datastore
+    LOGGER.debug("Loading datastore")
+    store = Store(store_location=store_location)
+
+    try:
+        # Process tables
+        source_ids, table_ingestion_id_mapping = process_tables(
+            store, kusto_client, ingest_client, database, operation_tag
+        )
+
+        LOGGER.info("Store data was sent for ADX ingestion")
+
+        # Monitor ingestion if wait is True
+        has_failures = False
+        if wait and source_ids:
+            has_failures = monitor_ingestion(ingest_client, source_ids, table_ingestion_id_mapping)
+
+        # Handle failures
+        should_abort = handle_failures(kusto_client, database, operation_tag, has_failures)
+        if should_abort:
+            return False
+
+        return True
+
+    except Exception as e:
+        LOGGER.exception("Error during ingestion process")
+        # Perform rollback using the tag
+        LOGGER.warning(f"Dropping data with tag: {operation_tag}")
+        _drop_by_tag(kusto_client, database, operation_tag)
+        raise e
+
+
 def dump_store_to_adx(
     store_folder: str,
     postgres_host: str,

cosmotech/coal/azure/adx/tables.py

@@ -5,8 +5,9 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
-from typing import Dict
+from typing import Dict, Any
 
+import pyarrow
 from azure.kusto.data import KustoClient
 
 from cosmotech.coal.utils.logger import LOGGER
@@ -39,6 +40,51 @@ def table_exists(client: KustoClient, database: str, table_name: str) -> bool:
     return False
 
 
+def check_and_create_table(kusto_client: KustoClient, database: str, table_name: str, data: pyarrow.Table) -> bool:
+    """
+    Check if a table exists and create it if it doesn't.
+
+    Args:
+        kusto_client: The Kusto client
+        database: The database name
+        table_name: The table name
+        data: The PyArrow table data
+
+    Returns:
+        bool: True if the table was created, False if it already existed
+    """
+    LOGGER.debug("  - Checking if table exists")
+    if not table_exists(kusto_client, database, table_name):
+        from cosmotech.coal.azure.adx.utils import create_column_mapping
+
+        mapping = create_column_mapping(data)
+        LOGGER.debug("  - Does not exist, creating it")
+        create_table(kusto_client, database, table_name, mapping)
+        return True
+    return False
+
+
+def _drop_by_tag(kusto_client: KustoClient, database: str, tag: str) -> None:
+    """
+    Drop all data with the specified tag.
+
+    Args:
+        kusto_client: The Kusto client
+        database: The database name
+        tag: The tag to drop data by
+    """
+    LOGGER.info(f"Dropping data with tag: {tag}")
+
+    try:
+        # Execute the drop by tag command
+        drop_command = f'.drop extents <| .show database extents where tags has "drop-by:{tag}"'
+        kusto_client.execute_mgmt(database, drop_command)
+        LOGGER.info("Drop by tag operation completed")
+    except Exception as e:
+        LOGGER.error(f"Error during drop by tag operation: {str(e)}")
+        LOGGER.exception("Drop by tag details")
+
+
 def create_table(client: KustoClient, database: str, table_name: str, schema: Dict[str, str]) -> bool:
     """
     Create a table in the database.