From 5b9c211a82bc12fc115c6dfd20951d0f88b7b1ed Mon Sep 17 00:00:00 2001 From: DAXNoobJustin Date: Sun, 12 Jan 2025 01:58:43 +0000 Subject: [PATCH 1/2] Adds a new function that allows users to query kql databases --- src/sempy_labs/_kql_databases.py | 38 ++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/src/sempy_labs/_kql_databases.py b/src/sempy_labs/_kql_databases.py index 3f462835..82523e84 100644 --- a/src/sempy_labs/_kql_databases.py +++ b/src/sempy_labs/_kql_databases.py @@ -139,3 +139,41 @@ def delete_kql_database(name: str, workspace: Optional[str | UUID] = None): print( f"{icons.green_dot} The '{name}' KQL database within the '{workspace_name}' workspace has been deleted." ) + +def query_kql_database(kql_database: [str | UUID], kql_query: str, workspace: Optional[str | UUID] = None) -> pd.DataFrame: + """ + Queries a KQL database. + + Parameters + ---------- + kql_database : str | uuid.UUID + The KQL database name or ID. + kql_query : str + The KQL query to execute. + workspace : str | uuid.UUID + The Fabric workspace name or ID. + + Returns + ------- + pandas.DataFrame + The result of the KQL query. + """ + + workspace_name, workspace_id = resolve_workspace_name_and_id(workspace) + kql_databases_df = labs.list_kql_databases(workspace_id) + + kql_database_search_attribute = 'KQL Database Id' if _is_valid_uuid(kql_database) else 'KQL Database Name' + kql_database_info = kql_databases_df.loc[kql_databases_df[kql_database_search_attribute] == kql_database] + kql_database_uri = kql_database_info['Query Service URI'].values[0] + kql_database_id = kql_database_info['KQL Database Id'].values[0] + + access_token = mssparkutils.credentials.getToken(kql_database_uri) + + kusto_query_result = spark.read \ + .format('com.microsoft.kusto.spark.synapse.datasource') \ + .option('accessToken', access_token) \ + .option('kustoCluster', kql_database_uri) \ + .option('kustoDatabase', kql_database_id) \ + .option('kustoQuery', kql_query).load() + + return kusto_query_result \ No newline at end of file From f51ef1e7ffca3eafd69a187036b55214e21b5d44 Mon Sep 17 00:00:00 2001 From: DAXNoobJustin Date: Sun, 12 Jan 2025 02:29:05 +0000 Subject: [PATCH 2/2] Add function to init and update return dataframe type --- src/sempy_labs/__init__.py | 1 + src/sempy_labs/_kql_databases.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/src/sempy_labs/__init__.py b/src/sempy_labs/__init__.py index 1803beca..67ea4249 100644 --- a/src/sempy_labs/__init__.py +++ b/src/sempy_labs/__init__.py @@ -413,6 +413,7 @@ "list_kql_databases", "create_kql_database", "delete_kql_database", + "query_kql_database", "create_eventhouse", "list_eventhouses", "delete_eventhouse", diff --git a/src/sempy_labs/_kql_databases.py b/src/sempy_labs/_kql_databases.py index 82523e84..36c4e4fa 100644 --- a/src/sempy_labs/_kql_databases.py +++ b/src/sempy_labs/_kql_databases.py @@ -176,4 +176,6 @@ def query_kql_database(kql_database: [str | UUID], kql_query: str, workspace: Op .option('kustoDatabase', kql_database_id) \ .option('kustoQuery', kql_query).load() + kusto_query_result = kusto_query_result.toPandas() + return kusto_query_result \ No newline at end of file