Updates for adx ingestion speed and on error rollback

Author: neomatamune

cosmotech/coal/azure/adx/ingestion.py

@@ -135,7 +135,6 @@ def check_ingestion_status(
     client: QueuedIngestClient,
     source_ids: List[str],
     timeout: Optional[int] = None,
-    logs: bool = False,
 ) -> Iterator[Tuple[str, IngestionStatus]]:
     """
     Check the status of ingestion operations.
@@ -144,7 +143,6 @@ def check_ingestion_status(
         client: The QueuedIngestClient to use
         source_ids: List of source IDs to check
         timeout: Timeout in seconds (default: 900)
-        logs: Whether to log detailed information
 
     Returns:
         Iterator of (source_id, status) tuples
@@ -185,7 +183,7 @@ def get_messages(queues):
 
     LOGGER.debug(T("coal.logs.adx.status_messages").format(success=len(successes), failure=len(failures)))
 
-    non_sent_ids = remaining_ids[:]
+    queued_ids = list(remaining_ids)
     # Process success and failure messages
     for messages, cast_func, status, log_function in [
         (successes, SuccessMessage, IngestionStatus.SUCCESS, LOGGER.debug),
@@ -207,11 +205,9 @@ def get_messages(queues):
             else:
                 # The message did not correspond to a known ID
                 continue
-            break
         else:
             # No message was found on the current list of messages for the given IDs
             continue
-        break
 
     # Check for timeouts
     actual_timeout = timeout if timeout is not None else default_timeout
@@ -221,7 +217,7 @@ def get_messages(queues):
             LOGGER.warning(T("coal.logs.adx.ingestion_timeout").format(source_id=source_id))
 
     # Yield results for remaining IDs
-    for source_id in non_sent_ids:
+    for source_id in queued_ids:
         yield source_id, _ingest_status[source_id]
 
 

cosmotech/coal/azure/adx/store.py

@@ -40,6 +40,7 @@ def send_pyarrow_table_to_adx(
         data_format=DataFormat.CSV,
         drop_by_tags=drop_by_tags,
         report_level=ReportLevel.FailuresAndSuccesses,
+        flush_immediately=True,
     )
 
     file_name = f"adx_{database}_{table_name}_{int(time.time())}_{uuid.uuid4()}.csv"

cosmotech/coal/azure/adx/wrapper.py

@@ -112,20 +112,19 @@ def ingest_dataframe(self, table_name: str, dataframe: Any, drop_by_tag: str = N
         return ingest_dataframe(self.ingest_client, self.database, table_name, dataframe, drop_by_tag)
 
     def check_ingestion_status(
-        self, source_ids: List[str], timeout: int = None, logs: bool = False
+        self, source_ids: List[str], timeout: int = None
     ) -> Iterator[Tuple[str, IngestionStatus]]:
         """
         Check the status of ingestion operations.
 
         Args:
             source_ids: List of source IDs to check
             timeout: Timeout in seconds (default: self.timeout)
-            logs: Whether to log detailed information
 
         Returns:
             Iterator of (source_id, status) tuples
         """
-        return check_ingestion_status(self.ingest_client, source_ids, timeout or self.timeout, logs)
+        return check_ingestion_status(self.ingest_client, source_ids, timeout or self.timeout)
 
     def _clear_ingestion_status_queues(self, confirmation: bool = False):
         """

cosmotech/coal/csm/engine/__init__.py

@@ -44,4 +44,4 @@ def apply_simple_csv_parameter_to_simulator(
         raise ValueError(f"Parameter {parameter_name} does not exists.")
 
 
-__all__ = [apply_simple_csv_parameter_to_simulator]
+__all__ = ["apply_simple_csv_parameter_to_simulator"]

cosmotech/coal/postgresql/runner.py

@@ -54,7 +54,7 @@ def send_runner_metadata_to_postgresql(
 
     # Generate PostgreSQL URI
     postgresql_full_uri = generate_postgresql_full_uri(
-        postgres_host, postgres_port, postgres_db, postgres_user, postgres_password
+        postgres_host, str(postgres_port), postgres_db, postgres_user, postgres_password
     )
 
     # Connect to PostgreSQL and update runner metadata
@@ -76,10 +76,10 @@ def send_runner_metadata_to_postgresql(
                   DO
                     UPDATE SET name = EXCLUDED.name, last_run_id = EXCLUDED.last_run_id;
             """
-            LOGGER.info(f"creating table {schema_table}")
+            LOGGER.info(T("coal.logs.postgreql.runner.creating_table").format(schema_table=schema_table))
             curs.execute(sql_create_table)
             conn.commit()
-            LOGGER.info(f"adding/updating runner metadata")
+            LOGGER.info(T("coal.logs.postgreql.runner.metadata"))
             curs.execute(
                 sql_upsert,
                 (
@@ -90,4 +90,4 @@ def send_runner_metadata_to_postgresql(
                 ),
             )
             conn.commit()
-            LOGGER.info("Runner metadata table has been updated")
+            LOGGER.info(T("coal.logs.postgreql.runner.metadata_updated"))

cosmotech/csm_data/commands/adx_send_data.py

@@ -45,11 +45,19 @@
     show_default=True,
     help="Wait for ingestion to complete",
 )
+@click.option(
+    "--tag",
+    envvar="CSM_DATA_ADX_TAG",
+    show_envvar=True,
+    default=None,
+    help="Optional tag to use for tracking and potential rollback of this ingestion operation",
+)
 def adx_send_data(
     adx_uri: str,
     adx_ingest_uri: str,
     database_name: str,
     wait: bool,
+    tag: str = None,
 ):
     # Import the function at the start of the command
     from cosmotech.coal.azure.adx.auth import create_ingest_client, create_kusto_client
@@ -61,8 +69,13 @@ def adx_send_data(
     from cosmotech.coal.azure.adx import type_mapping
 
     import time
+    import uuid
     from cosmotech.coal.azure.adx import IngestionStatus
 
+    # Generate operation tag if not provided
+    operation_tag = tag or f"op-{str(uuid.uuid4())}"
+    LOGGER.debug(f"Starting ingestion operation with tag: {operation_tag}")
+
     LOGGER.debug("Initializing clients")
     kusto_client = create_kusto_client(adx_uri)
     ingest_client = create_ingest_client(adx_ingest_uri)
@@ -72,14 +85,14 @@ def adx_send_data(
     s = Store()
     source_ids = []
     LOGGER.debug("Listing tables")
-    table_list = list(s.list_tables())[:3]
+    table_list = list(s.list_tables())
     table_ingestion_id_mapping = dict()
     for target_table_name in table_list:
         LOGGER.info(f"Working on table: {target_table_name}")
         data = s.get_table(target_table_name)
 
         if data.num_rows < 1:
-            LOGGER.warn(f"Table {target_table_name} has no rows - skipping it")
+            LOGGER.warning(f"Table {target_table_name} has no rows - skipping it")
             continue
 
         LOGGER.debug("  - Checking if table exists")
@@ -99,38 +112,90 @@ def adx_send_data(
             create_table(kusto_client, database, target_table_name, mapping)
 
         LOGGER.debug(f"Sending data to the table {target_table_name}")
-        result = send_pyarrow_table_to_adx(ingest_client, database, target_table_name, data, None)
+        # Use the operation_tag as the drop_by_tag parameter
+        result = send_pyarrow_table_to_adx(ingest_client, database, target_table_name, data, operation_tag)
         source_ids.append(result.source_id)
         table_ingestion_id_mapping[result.source_id] = target_table_name
 
+    # Track if any failures occur
+    has_failures = False
+
     LOGGER.info("Store data was sent for ADX ingestion")
-    if wait:
-        LOGGER.info("Waiting for ingestion of data to finish")
-        import tqdm
-
-        with tqdm.tqdm(desc="Ingestion status", total=len(source_ids)) as pbar:
-            while any(
-                map(
-                    lambda _status: _status[1] in (IngestionStatus.QUEUED, IngestionStatus.UNKNOWN),
-                    results := list(check_ingestion_status(ingest_client, source_ids)),
-                )
-            ):
-                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.remove(ingestion_id)
-
-                if os.environ.get("CSM_USE_RICH", "False").lower() in ("true", "1", "yes", "t", "y"):
-                    for _ in range(10):
-                        time.sleep(1)
-                        pbar.update(0)
+    try:
+        if wait:
+            LOGGER.info("Waiting for ingestion of data to finish")
+            import tqdm
+
+            with tqdm.tqdm(desc="Ingestion status", total=len(source_ids)) as pbar:
+                while any(
+                    list(
+                        map(
+                            lambda _status: _status[1] in (IngestionStatus.QUEUED, IngestionStatus.UNKNOWN),
+                            results := list(check_ingestion_status(ingest_client, source_ids)),
+                        )
+                    )
+                ):
+                    # 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.remove(ingestion_id)
+
+                    time.sleep(1)
+                    if os.environ.get("CSM_USE_RICH", "False").lower() in ("true", "1", "yes", "t", "y"):
+                        pbar.refresh()
                 else:
-                    time.sleep(10)
-            pbar.update(len(source_ids))
-        LOGGER.info("All data got ingested")
+                    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))
+            LOGGER.info("All data ingestion attempts completed")
+
+        # If any ingestion failed, perform rollback
+        if has_failures:
+            LOGGER.warning(f"Failures detected during ingestion - dropping data with tag: {operation_tag}")
+            _drop_by_tag(kusto_client, database, operation_tag)
+
+    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
+
+    if has_failures:
+        click.Abort()
+
+
+def _drop_by_tag(kusto_client, database, tag):
+    """
+    Drop all data with the specified tag
+    """
+    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")
 
 
 if __name__ == "__main__":

tests/unit/coal/test_azure/test_adx/test_adx_ingestion.py

@@ -357,31 +357,6 @@ def test_check_ingestion_status_with_timeout(
         assert len(result) == 1
         assert result[0] == (source_id, IngestionStatus.TIMEOUT)
 
-    @patch("cosmotech.coal.azure.adx.ingestion.KustoIngestStatusQueues")
-    def test_check_ingestion_status_with_logs(self, mock_status_queues_class, mock_ingest_client, mock_status_queues):
-        """Test the check_ingestion_status function with logs enabled."""
-        # Arrange
-        source_id = "source-id-logs"
-        _ingest_status[source_id] = IngestionStatus.QUEUED
-        _ingest_times[source_id] = time.time()
-
-        # Set up mock status queues with empty queues
-        mock_status_queues_class.return_value = mock_status_queues
-        mock_success_queue = MagicMock()
-        mock_success_queue.receive_messages.return_value = []
-        mock_status_queues.success._get_queues.return_value = [mock_success_queue]
-        mock_failure_queue = MagicMock()
-        mock_failure_queue.receive_messages.return_value = []
-        mock_status_queues.failure._get_queues.return_value = [mock_failure_queue]
-
-        # Act
-        result = list(check_ingestion_status(mock_ingest_client, [source_id], logs=True))
-
-        # Assert
-        assert len(result) == 1
-        # The status should still be QUEUED since no messages were found and no timeout occurred
-        assert result[0] == (source_id, IngestionStatus.QUEUED)
-
     @patch("cosmotech.coal.azure.adx.ingestion.KustoIngestStatusQueues")
     def test_check_ingestion_status_unknown_id(self, mock_status_queues_class, mock_ingest_client, mock_status_queues):
         """Test the check_ingestion_status function with an unknown source ID."""