cosmotech.coal.azure

cosmotech.coal.azure.adx.auth

Functions:

Name	Description
create_kusto_client	Create a KustoClient for querying ADX.
create_ingest_client	Create a QueuedIngestClient for ingesting data to ADX.
initialize_clients	Initialize and return the Kusto and ingest clients.
get_cluster_urls	Generate cluster and ingest URLs from cluster name and region.

create_kusto_client(cluster_url, client_id=None, client_secret=None, tenant_id=None)

Create a KustoClient for querying ADX.

Parameters:

Name	Type	Description	Default
cluster_url	str	The URL of the ADX cluster	required
client_id	Optional[str]	Azure client ID (optional, will use environment variable if not provided)	None
client_secret	Optional[str]	Azure client secret (optional, will use environment variable if not provided)	None
tenant_id	Optional[str]	Azure tenant ID (optional, will use environment variable if not provided)	None

Returns:

Name	Type	Description
KustoClient	KustoClient	A client for querying ADX

create_ingest_client(ingest_url, client_id=None, client_secret=None, tenant_id=None)

Create a QueuedIngestClient for ingesting data to ADX.

Parameters:

Name	Type	Description	Default
ingest_url	str	The ingestion URL of the ADX cluster	required
client_id	Optional[str]	Azure client ID (optional, will use environment variable if not provided)	None
client_secret	Optional[str]	Azure client secret (optional, will use environment variable if not provided)	None
tenant_id	Optional[str]	Azure tenant ID (optional, will use environment variable if not provided)	None

Returns:

Name	Type	Description
QueuedIngestClient	QueuedIngestClient	A client for ingesting data to ADX

initialize_clients(adx_uri, adx_ingest_uri)

Initialize and return the Kusto and ingest clients.

Parameters:

Name	Type	Description	Default
adx_uri	str	The Azure Data Explorer resource URI	required
adx_ingest_uri	str	The Azure Data Explorer resource ingest URI	required

Returns:

Name	Type	Description
tuple	Tuple[KustoClient, QueuedIngestClient]	(kusto_client, ingest_client)

get_cluster_urls(cluster_name, cluster_region)

Generate cluster and ingest URLs from cluster name and region.

Parameters:

Name	Type	Description	Default
cluster_name	str	The name of the ADX cluster	required
cluster_region	str	The region of the ADX cluster	required

Returns:

Name	Type	Description
tuple	Tuple[str, str]	(cluster_url, ingest_url)

cosmotech.coal.azure.adx.ingestion

Functions:

Name	Description
ingest_dataframe	Ingest a pandas DataFrame into an ADX table.
send_to_adx	Send a list of dictionaries to an ADX table.
check_ingestion_status	Check the status of ingestion operations.
monitor_ingestion	Monitor the ingestion process with progress reporting.
handle_failures	Handle any failures and perform rollbacks if needed.
clear_ingestion_status_queues	Clear all data in the ingestion status queues.

ingest_dataframe(client, database, table_name, dataframe, drop_by_tag=None)

Ingest a pandas DataFrame into an ADX table.

Parameters:

Name	Type	Description	Default
client	QueuedIngestClient	The QueuedIngestClient to use	required
database	str	The name of the database	required
table_name	str	The name of the table	required
dataframe	DataFrame	The DataFrame to ingest	required
drop_by_tag	Optional[str]	Tag used for the drop by capacity of the Cosmotech API	None

Returns:

Type	Description
	The ingestion result with source_id for status tracking

send_to_adx(query_client, ingest_client, database, dict_list, table_name, ignore_table_creation=True, drop_by_tag=None)

Send a list of dictionaries to an ADX table.

Parameters:

Name	Type	Description	Default
query_client	KustoClient	The KustoClient for querying	required
ingest_client	QueuedIngestClient	The QueuedIngestClient for ingestion	required
database	str	The name of the database	required
dict_list	List[Dict]	The list of dictionaries to send	required
table_name	str	The name of the table	required
ignore_table_creation	bool	If False, will create the table if it doesn't exist	True
drop_by_tag	Optional[str]	Tag used for the drop by capacity of the Cosmotech API	None

Returns:

Type	Description
	The ingestion result with source_id for status tracking

check_ingestion_status(client, source_ids, timeout=None)

Check the status of ingestion operations.

Parameters:

Name	Type	Description	Default
client	QueuedIngestClient	The QueuedIngestClient to use	required
source_ids	List[str]	List of source IDs to check	required
timeout	Optional[int]	Timeout in seconds (default: 900)	None

Returns:

Type	Description
Iterator[Tuple[str, IngestionStatus]]	Iterator of (source_id, status) tuples

monitor_ingestion(ingest_client, source_ids, table_ingestion_id_mapping)

Monitor the ingestion process with progress reporting.

Parameters:

Name	Type	Description	Default
ingest_client	QueuedIngestClient	The ingest client	required
source_ids	List[str]	List of source IDs to monitor	required
table_ingestion_id_mapping	Dict[str, str]	Mapping of source IDs to table names	required

Returns:

Name	Type	Description
bool	bool	True if any failures occurred, False otherwise

handle_failures(kusto_client, database, operation_tag, has_failures)

Handle any failures and perform rollbacks if needed.

Parameters:

Name	Type	Description	Default
kusto_client	KustoClient	The Kusto client	required
database	str	The database name	required
operation_tag	str	The operation tag for tracking	required
has_failures	bool	Whether any failures occurred	required

Returns:

Name	Type	Description
bool	bool	True if the process should abort, False otherwise

clear_ingestion_status_queues(client, confirmation=False)

Clear all data in the ingestion status queues. DANGEROUS: This will clear all queues for the entire ADX cluster.

Parameters:

Name	Type	Description	Default
client	QueuedIngestClient	The QueuedIngestClient to use	required
confirmation	bool	Must be True to proceed with clearing	False

cosmotech.coal.azure.adx.query

Functions:

Name	Description
run_query	Execute a simple query on the database.
run_command_query	Execute a command query on the database.

run_query(client, database, query)

Execute a simple query on the database.

Parameters:

Name	Type	Description	Default
client	KustoClient	The KustoClient to use	required
database	str	The name of the database	required
query	str	The query to execute	required

Returns:

Name	Type	Description
KustoResponseDataSet	KustoResponseDataSet	The results of the query

run_command_query(client, database, query)

Execute a command query on the database.

Parameters:

Name	Type	Description	Default
client	KustoClient	The KustoClient to use	required
database	str	The name of the database	required
query	str	The query to execute	required

Returns:

Name	Type	Description
KustoResponseDataSet	KustoResponseDataSet	The results of the query

cosmotech.coal.azure.adx.runner

ADX runner data ingestion module.

This module provides functions for ingesting runner data into Azure Data Explorer.

Functions:

Name	Description
prepare_csv_content	Navigate through folder_path to generate csv information for each csv file in it.
construct_create_query	Construct ADX table creation queries for the given CSV files.
insert_csv_files	Insert CSV files into ADX tables.
send_runner_data	Send runner data to ADX.

prepare_csv_content(folder_path)

Navigate through folder_path to generate csv information for each csv file in it.

Parameters:

Name	Type	Description	Default
folder_path	str	Path to the folder containing CSV files	required

Returns:

Name	Type	Description
	Dict[str, Dict[str, Any]]	A map of filename to file_infos
	Dict[str, Dict[str, Any]]	file infos:
dict	Dict[str, Dict[str, Any]]	filename -> filename as a string without path & extension headers -> map of column_name -> column_type

construct_create_query(files_data)

Construct ADX table creation queries for the given CSV files.

Parameters:

Name	Type	Description	Default
files_data	Dict[str, Dict[str, Any]]	Map of filename to file_infos as returned by prepare_csv_content	required

Returns:

Type	Description
Dict[str, str]	Map of table_name to creation query

insert_csv_files(files_data, ingest_client, runner_id, database, wait=False, wait_limit=5, wait_duration=8)

Insert CSV files into ADX tables.

Parameters:

Name	Type	Description	Default
files_data	Dict[str, Dict[str, Any]]	Map of filename to file_infos as returned by prepare_csv_content	required
kusto_client		The KustoClient for querying	required
ingest_client	QueuedIngestClient	The QueuedIngestClient for ingestion	required
runner_id	str	Runner ID to use as a tag	required
database	str	ADX database name	required
wait	bool	Whether to wait for ingestion to complete	False
wait_limit	int	Number of retries while waiting	5
wait_duration	int	Duration between each try while waiting	8

send_runner_data(dataset_absolute_path, parameters_absolute_path, runner_id, adx_uri, adx_ingest_uri, database_name, send_parameters=False, send_datasets=False, wait=False)

Send runner data to ADX.

Parameters:

Name	Type	Description	Default
dataset_absolute_path	str	Path to the dataset folder	required
parameters_absolute_path	str	Path to the parameters folder	required
runner_id	str	Runner ID to use as a tag	required
adx_uri	str	ADX cluster URI	required
adx_ingest_uri	str	ADX ingestion URI	required
database_name	str	ADX database name	required
send_parameters	bool	Whether to send parameters	False
send_datasets	bool	Whether to send datasets	False
wait	bool	Whether to wait for ingestion to complete	False

cosmotech.coal.azure.adx.store

Functions:

Name	Description
send_table_data	Send a PyArrow table to ADX.
process_tables	Process all tables in the store.
send_pyarrow_table_to_adx
send_store_to_adx	Send data from the store to Azure Data Explorer.

send_table_data(ingest_client, database, table_name, data, operation_tag)

Send a PyArrow table to ADX.

Parameters:

Name	Type	Description	Default
ingest_client	QueuedIngestClient	The ingest client	required
database	str	The database name	required
table_name	str	The table name	required
data	Table	The PyArrow table data	required
operation_tag	str	The operation tag for tracking	required

Returns:

Name	Type	Description
tuple	Tuple[str, str]	(source_id, table_name)

process_tables(store, kusto_client, ingest_client, database, operation_tag)

Process all tables in the store.

Parameters:

Name	Type	Description	Default
store	Store	The data store	required
kusto_client	KustoClient	The Kusto client	required
ingest_client	QueuedIngestClient	The ingest client	required
database	str	The database name	required
operation_tag	str	The operation tag for tracking	required

Returns:

Name	Type	Description
tuple	Tuple[List[str], Dict[str, str]]	(source_ids, table_ingestion_id_mapping)

send_pyarrow_table_to_adx(client, database, table_name, table_data, drop_by_tag=None)

send_store_to_adx(adx_uri, adx_ingest_uri, database_name, wait=False, tag=None, store_location=None)

Send data from the store to Azure Data Explorer.

Parameters:

Name	Type	Description	Default
adx_uri	str	The Azure Data Explorer resource URI	required
adx_ingest_uri	str	The Azure Data Explorer resource ingest URI	required
database_name	str	The database name	required
wait	bool	Whether to wait for ingestion to complete	False
tag	Optional[str]	The operation tag for tracking (will generate a unique one if not provided)	None
store_location	Optional[str]	Optional store location (uses default if not provided)	None

Returns:

Name	Type	Description
bool	Union[bool, Any]	True if successful, False otherwise

cosmotech.coal.azure.adx.tables

Functions:

Name	Description
table_exists	Check if a table exists in the database.
check_and_create_table	Check if a table exists and create it if it doesn't.
create_table	Create a table in the database.

table_exists(client, database, table_name)

Check if a table exists in the database.

Parameters:

Name	Type	Description	Default
client	KustoClient	The KustoClient to use	required
database	str	The name of the database	required
table_name	str	The name of the table to check	required

Returns:

Name	Type	Description
bool	bool	True if the table exists, False otherwise

check_and_create_table(kusto_client, database, table_name, data)

Check if a table exists and create it if it doesn't.

Parameters:

Name	Type	Description	Default
kusto_client	KustoClient	The Kusto client	required
database	str	The database name	required
table_name	str	The table name	required
data	Table	The PyArrow table data	required

Returns:

Name	Type	Description
bool	bool	True if the table was created, False if it already existed

create_table(client, database, table_name, schema)

Create a table in the database.

Parameters:

Name	Type	Description	Default
client	KustoClient	The KustoClient to use	required
database	str	The name of the database	required
table_name	str	The name of the table to create	required
schema	Dict[str, str]	Dictionary mapping column names to ADX types	required

Returns:

Name	Type	Description
bool	bool	True if the table was created successfully, False otherwise

cosmotech.coal.azure.adx.utils

Functions:

Name	Description
create_column_mapping	Create a column mapping for a PyArrow table.
type_mapping	Map Python types to ADX types.

create_column_mapping(data)

Create a column mapping for a PyArrow table.

Parameters:

Name	Type	Description	Default
data	Table	The PyArrow table data	required

Returns:

Name	Type	Description
dict	Dict[str, str]	A mapping of column names to their ADX types

type_mapping(key, key_example_value)

Map Python types to ADX types.

Parameters:

Name	Type	Description	Default
key	str	The name of the key	required
key_example_value	Any	A possible value of the key	required

Returns:

Name	Type	Description
str	str	The name of the type used in ADX

cosmotech.coal.azure.blob

Azure Blob Storage operations module.

This module provides functions for interacting with Azure Blob Storage, including uploading data from the Store.

Functions:

Name	Description
dump_store_to_azure	Dump Store data to Azure Blob Storage.
delete_azure_blobs

dump_store_to_azure(configuration=Configuration(), selected_tables=[])

Dump Store data to Azure Blob Storage.

Parameters:

Name	Type	Description	Default
configuration	Configuration	Configuration utils class	Configuration()
selected_tables	list[str]	List of tables name	[]

Raises:

Type	Description
ValueError	If the output type is invalid

delete_azure_blobs(configuration=Configuration())

cosmotech.coal.azure.storage

Azure Storage operations module.

This module provides functions for interacting with Azure Storage, including uploading files to blob storage.

Functions:

Name	Description
upload_file	Upload a single file to Azure Blob Storage.
upload_folder	Upload files from a folder to Azure Blob Storage.

upload_file(file_path, blob_name, az_storage_sas_url, file_prefix='')

Upload a single file to Azure Blob Storage.

Parameters:

Name	Type	Description	Default
file_path	Path	Path to the file to upload	required
blob_name	str	Name of the blob container	required
az_storage_sas_url	str	SAS URL for the Azure Storage account	required
file_prefix	str	Prefix to add to the file name in the blob	''

upload_folder(source_folder, blob_name, az_storage_sas_url, file_prefix='', recursive=False)

Upload files from a folder to Azure Blob Storage.

Parameters:

Name	Type	Description	Default
source_folder	str	Path to the folder containing files to upload	required
blob_name	str	Name of the blob container	required
az_storage_sas_url	str	SAS URL for the Azure Storage account	required
file_prefix	str	Prefix to add to the file names in the blob	''
recursive	bool	Whether to recursively upload files from subdirectories	False