def test_null_values_in_data(self, mock_post): """Tests response with null values in non nullable column types""" client = KustoClient(self.HOST) query = "PrimaryResultName" response = client.execute_query("PythonTest", query) assert response is not None
def test_identifiable_os(self, mock_post): """Tests identifiable OS doesn't fail when composing its socket options""" with patch.object(sys, "platform", "win32"): client = KustoClient("https://somecluster.kusto.windows.net") query = """print dynamic(123)""" row = client.execute_query("PythonTest", query).primary_results[0].rows[0] assert isinstance(row[0], int)
def setup_class(cls): cls.engine_cs = os.environ.get("ENGINE_CONNECTION_STRING") or "" cls.ai_engine_cs = os.environ.get("APPLICATION_INSIGHTS_ENGINE_CONNECTION_STRING") or "" cls.app_id = os.environ.get("APP_ID") cls.app_key = os.environ.get("APP_KEY") cls.auth_id = os.environ.get("AUTH_ID") cls.test_db = os.environ.get("TEST_DATABASE") cls.ai_test_db = os.environ.get("APPLICATION_INSIGHTS_TEST_DATABASE") # name of e2e database could be changed if not all([cls.engine_cs, cls.test_db, cls.ai_engine_cs, cls.ai_test_db]): pytest.skip("E2E environment is missing") # Init clients cls.streaming_test_table = "BigChunkus" cls.streaming_test_table_query = cls.streaming_test_table + " | order by timestamp" cls.ai_test_table_cmd = ".show tables" cls.client = KustoClient(cls.engine_kcsb_from_env()) cls.ai_client = KustoClient(cls.engine_kcsb_from_env(True)) cls.input_folder_path = cls.get_file_path() with open(os.path.join(cls.input_folder_path, "big.json")) as f: cls.test_streaming_data = json.load(f)
def test_null_values_in_data(self, mock_post): """Tests response with null values in non nullable column types""" client = KustoClient("https://somecluster.kusto.windows.net") query = "PrimaryResultName" response = client.execute_query("PythonTest", query) assert response is not None
def test_dynamic(self, mock_post): """Tests dynamic responses.""" client = KustoClient(self.HOST) query = """print dynamic(123), dynamic("123"), dynamic("test bad json"),""" """ dynamic(null), dynamic('{"rowId":2,"arr":[0,2]}'), dynamic({"rowId":2,"arr":[0,2]})""" row = client.execute_query("PythonTest", query).primary_results[0].rows[0] self._assert_dynamic_response(row)
class KustoStreamingIngestClient(BaseIngestClient): """Kusto streaming ingest client for Python. KustoStreamingIngestClient works with both 2.x and 3.x flavors of Python. All primitive types are supported. Tests are run using pytest. """ def __init__(self, kcsb: Union[KustoConnectionStringBuilder, str]): """Kusto Streaming Ingest Client constructor. :param KustoConnectionStringBuilder kcsb: The connection string to initialize KustoClient. """ self._kusto_client = KustoClient(kcsb) def set_proxy(self, proxy_url: str): self._kusto_client.set_proxy(proxy_url) def ingest_from_file(self, file_descriptor: Union[FileDescriptor, str], ingestion_properties: IngestionProperties) -> IngestionResult: """Ingest from local files. :param file_descriptor: a FileDescriptor to be ingested. :param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties. """ stream_descriptor = StreamDescriptor.from_file_descriptor(file_descriptor) with stream_descriptor.stream: return self.ingest_from_stream(stream_descriptor, ingestion_properties) def ingest_from_stream(self, stream_descriptor: Union[StreamDescriptor, IO[AnyStr]], ingestion_properties: IngestionProperties) -> IngestionResult: """Ingest from io streams. :param azure.kusto.ingest.StreamDescriptor stream_descriptor: An object that contains a description of the stream to be ingested. :param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties. """ return self._ingest_from_stream_with_client_request_id(stream_descriptor, ingestion_properties, None) def _ingest_from_stream_with_client_request_id( self, stream_descriptor: Union[StreamDescriptor, IO[AnyStr]], ingestion_properties: IngestionProperties, client_request_id: Optional[str] ) -> IngestionResult: stream_descriptor = BaseIngestClient._prepare_stream(stream_descriptor, ingestion_properties) additional_properties = None if client_request_id: additional_properties = ClientRequestProperties() additional_properties.client_request_id = client_request_id self._kusto_client.execute_streaming_ingest( ingestion_properties.database, ingestion_properties.table, stream_descriptor.stream, ingestion_properties.format.name, additional_properties, mapping_name=ingestion_properties.ingestion_mapping_reference, ) return IngestionResult(IngestionStatus.SUCCESS, ingestion_properties.database, ingestion_properties.table, stream_descriptor.source_id)
def test_admin_then_query(self, mock_post): """Tests admin then query.""" client = KustoClient("https://somecluster.kusto.windows.net") query = ".show tables | project DatabaseName, TableName" response = client.execute_mgmt("PythonTest", query) assert response.errors_count == 0 assert len(response) == 4 results = list(response.primary_results[0]) assert len(results) == 2 assert response[0].table_kind == WellKnownDataSet.PrimaryResult assert response[1].table_kind == WellKnownDataSet.QueryProperties assert response[2].table_kind == WellKnownDataSet.QueryCompletionInformation assert response[3].table_kind == WellKnownDataSet.TableOfContents
def test_sanity_control_command(self, mock_post): """Tests contol command.""" client = KustoClient("https://somecluster.kusto.windows.net") response = client.execute_mgmt("NetDefaultDB", ".show version") assert len(response) == 1 primary_table = response.primary_results[0] row_count = 0 for _ in primary_table: row_count += 1 assert row_count == 1 result = primary_table[0] assert result["BuildVersion"] == "1.0.6693.14577" assert result["BuildTime"] == datetime(year=2018, month=4, day=29, hour=8, minute=5, second=54, tzinfo=UTC) assert result["ServiceType"] == "Engine" assert result["ProductVersion"] == "KustoMain_2018.04.29.5"
def test_partial_results(self, mock_post): """Tests partial results.""" client = KustoClient("https://somecluster.kusto.windows.net") query = """set truncationmaxrecords = 5; range x from 1 to 10 step 1""" properties = ClientRequestProperties() properties.set_option(ClientRequestProperties.results_defer_partial_query_failures_option_name, False) self.assertRaises(KustoServiceError, client.execute_query, "PythonTest", query, properties) properties.set_option(ClientRequestProperties.results_defer_partial_query_failures_option_name, True) response = client.execute_query("PythonTest", query, properties) assert response.errors_count == 1 assert "E_QUERY_RESULT_SET_TOO_LARGE" in response.get_exceptions()[0] assert len(response) == 3 results = list(response.primary_results[0]) assert len(results) == 5 assert results[0]["x"] == 1
def test_partial_results(self, mock_post): """Tests partial results.""" client = KustoClient(self.HOST) query = """set truncationmaxrecords = 5; range x from 1 to 10 step 1""" properties = ClientRequestProperties() properties.set_option( ClientRequestProperties. results_defer_partial_query_failures_option_name, False) self.assertRaises(KustoServiceError, client.execute_query, "PythonTest", query, properties) properties.set_option( ClientRequestProperties. results_defer_partial_query_failures_option_name, True) response = client.execute_query("PythonTest", query, properties) self._assert_partial_results_response(response)
def authenticate_to_kusto(cluster): """Authenticate and return kusto connection client""" kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication( cluster, CLIENT_ID, CLIENT_SECRET, AUTHORITY_ID) # The authentication method will be taken from the chosen KustoConnectionStringBuilder. kusto_client = KustoClient(kcsb) return kusto_client
def test_null_values_in_data(self, mock_post, method): """Tests response with null values in non nullable column types""" client = KustoClient(self.HOST) query = "PrimaryResultName" response = method.__call__(client, "PythonTest", query) assert response is not None
def test_sanity_data_frame(self, mock_post, method): """Tests KustoResponse to pandas.DataFrame.""" client = KustoClient(self.HOST) response = method.__call__(client, "PythonTest", "Deft") data_frame = dataframe_from_result_table( get_response_first_primary_result(response)) self._assert_sanity_data_frame_response(data_frame)
def test_partial_results(self, mock_post, method): """Tests partial results.""" client = KustoClient(self.HOST) query = """set truncationmaxrecords = 5; range x from 1 to 10 step 1""" properties = ClientRequestProperties() properties.set_option( ClientRequestProperties. results_defer_partial_query_failures_option_name, False) with pytest.raises(KustoMultiApiError) as e: response = method.__call__(client, "PythonTest", query, properties=properties) if type(response) == KustoStreamingResponseDataSet: results = list(get_response_first_primary_result(response)) errors = e.value.get_api_errors() assert len(errors) == 1 assert errors[0].code == "LimitsExceeded" properties.set_option( ClientRequestProperties. results_defer_partial_query_failures_option_name, True) response = method.__call__(client, "PythonTest", query, properties=properties) self._assert_partial_results_response(response)
def test_partial_results(self, mock_post, mock_aad): """ Tests partial results. """ client = KustoClient("https://somecluster.kusto.windows.net") query = """\ set truncationmaxrecords = 1; range x from 1 to 2 step 1 """ self.assertRaises(KustoServiceError, client.execute_query, "PythonTest", query) response = client.execute_query("PythonTest", query, accept_partial_results=True) self.assertTrue(response.has_exceptions()) self.assertEqual(response.get_exceptions()[0]['error']['code'], 'LimitsExceeded') self.assertEqual(response.get_table_count(), 5) results = list(response.iter_all()) self.assertEqual(len(results), 1) self.assertEqual(results[0]['x'], 1)
def test_dynamic(self, mock_post, method): """Tests dynamic responses.""" client = KustoClient(self.HOST) query = """print dynamic(123), dynamic("123"), dynamic("test bad json"),""" """ dynamic(null), dynamic('{"rowId":2,"arr":[0,2]}'), dynamic({"rowId":2,"arr":[0,2]})""" row = get_table_first_row( get_response_first_primary_result( method.__call__(client, "PythonTest", query))) self._assert_dynamic_response(row)
def test_identifiable_os(self, mock_post, method): """Tests identifiable OS doesn't fail when composing its socket options""" with patch.object(sys, "platform", "win32"): client = KustoClient("https://somecluster.kusto.windows.net") query = """print dynamic(123)""" row = get_table_first_row( get_response_first_primary_result( method.__call__(client, "PythonTest", query))) assert isinstance(row[0], int)
def __init__(self, kcsb: Union[str, KustoConnectionStringBuilder]): """Kusto Ingest Client constructor. :param kcsb: The connection string to initialize KustoClient. """ if not isinstance(kcsb, KustoConnectionStringBuilder): kcsb = KustoConnectionStringBuilder(kcsb) self._connection_datasource = kcsb.data_source self._resource_manager = _ResourceManager(KustoClient(kcsb)) self._endpoint_service_type = None self._suggested_endpoint_uri = None
def get_all_db_table_records_count(kusto_client: KustoClient, database: str, table_prefix: str = ''): # Count Records from CO Table and TEMP Tables in All Databases query_filter = f"union database(\"*\").CO*|union database(\"*\").TEMP*|count" print(f"query on database: {database} by query_filter: {query_filter}") response = kusto_client.execute(database, query_filter) result = response.primary_results[0] records_count = result.to_dict().get('data')[0] return records_count.get('Count')
def _get_client_for_cluster(cluster: str) -> KustoClient: # If we call 'with_az_cli_authentication' directly, in case of failure we will get an un-informative exception. # As a workaround, we first attempt to manually get the Azure CLI token, and see if it works. # Get rid of this workaround once this is resolved: https://github.com/Azure/azure-kusto-python/issues/240 stored_token = _get_azure_cli_auth_token() if stored_token is None: _logger.info("Failed to get Azure CLI token, falling back to AAD device authentication") connection_string_builder = KustoConnectionStringBuilder.with_aad_device_authentication(cluster) else: connection_string_builder = KustoConnectionStringBuilder.with_az_cli_authentication(cluster) return KustoClient(connection_string_builder)
def __init__(self, kusto_cluster, client_id=None, client_secret=None, username=None, password=None, authority=None): """ Kusto Client constructor. Parameters ---------- kusto_cluster : str Kusto cluster endpoint. Example: https://ingest-help.kusto.windows.net client_id : str The AAD application ID of the application making the request to Kusto client_secret : str The AAD application key of the application making the request to Kusto. if this is given, then username/password should not be. username : str The username of the user making the request to Kusto. if this is given, then password must follow and the client_secret should not be given. password : str The password matching the username of the user making the request to Kusto version : 'v1', optional REST API version, defaults to v1. authority : 'microsoft.com', optional In case your tenant is not microsoft please use this param. """ self._kusto_client = KustoClient(kusto_cluster, client_id=client_id, client_secret=client_secret, username=username, password=password, authority=authority) self._last_time_refreshed_containers = datetime.min self._temp_storage_objects = None self._last_time_refreshed_queues = datetime.min self._queues = None self._last_time_refreshed_token = datetime.min self._kusto_token = None
def test_dynamic(self, mock_post): """Tests dynamic responses.""" client = KustoClient("https://somecluster.kusto.windows.net") query = """print dynamic(123), dynamic("123"), dynamic("test bad json"),""" """ dynamic(null), dynamic('{"rowId":2,"arr":[0,2]}'), dynamic({"rowId":2,"arr":[0,2]})""" row = client.execute_query("PythonTest", query).primary_results[0].rows[0] assert isinstance(row[0], int) assert row[0] == 123 assert isinstance(row[1], str) assert row[1] == "123" assert isinstance(row[2], str) assert row[2] == "test bad json" assert row[3] is None assert isinstance(row[4], str) assert row[4] == '{"rowId":2,"arr":[0,2]}' assert isinstance(row[5], dict) assert row[5] == {"rowId": 2, "arr": [0, 2]}
def test_custom_request_id(self, mock_post, method): """Test query V2.""" client = KustoClient(self.HOST) properties = ClientRequestProperties() request_id = "test_request_id" properties.client_request_id = request_id response = method.__call__(client, "PythonTest", "Deft", properties=properties) self._assert_sanity_query_response(response) self._assert_client_request_id(mock_post.call_args.kwargs, value=request_id)
def test_sanity_control_command(self, mock_post, mock_aad): """ Tests contol command. """ client = KustoClient("https://somecluster.kusto.windows.net") response = client.execute_mgmt("NetDefaultDB", ".show version") self.assertEqual(response.get_table_count(), 1) row_count = 0 for _ in response.iter_all(): row_count += 1 self.assertEqual(row_count, 1) result = list(response.iter_all())[0] self.assertEqual(result['BuildVersion'], "1.0.6693.14577") self.assertEqual( result['BuildTime'], datetime(year=2018, month=4, day=29, hour=8, minute=5, second=54, tzinfo=tzutc())) self.assertEqual(result['ServiceType'], "Engine") self.assertEqual(result['ProductVersion'], "KustoMain_2018.04.29.5")
def test_proxy_token_providers(self, mock_get, proxy_kcsb): """Test query V2.""" proxy = "https://my_proxy.sample" kcsb, auth_supports_proxy = proxy_kcsb client = KustoClient(kcsb) client.set_proxy(proxy) assert client._proxy_url == proxy expected_dict = {"http": proxy, "https": proxy} if not auth_supports_proxy: return assert client._aad_helper.token_provider._proxy_dict == expected_dict assert client._session.proxies == expected_dict CloudSettings._cloud_cache.clear() client._aad_helper.token_provider._init_resources() mock_get.assert_called_with( "https://somecluster.kusto.windows.net/v1/rest/auth/metadata", proxies=expected_dict)
def __init__(self, cluster, database, table, clientId, clientSecret, authority="microsoft.com", resetTable=False): """ Parameters ---------- cluster : str Azure Data Explorer (ADX) cluster address. eg, 'CDOC.kusto.windows.net' database : str Azure Data Explorer (ADX) database name. eg, 'TestDb' table : str Azure Data Explorer (ADX) table name. eg, 'OutputTable' clientId : str Azure Data Explorer (ADX) client Id that has permissions to access ADX. clientSecret : str Azure Data Explorer (ADX) access key. Used along with client Id. authority : str Azure Data Explorer (ADX) authority. Optional. When not specified, 'microsoft.com' is used. resetTable : bool Default is False. If True, the existing data in the destination table is dropped before new data is logged. """ self.running = True self.batchSize = 10000 self.flushDuration = timedelta(milliseconds = 1000) self.lastUploadTime = datetime.utcnow() self.initTable = False self.nextBatch = list() self.currentBatch = None self.lock = threading.Lock() self.resetTable = resetTable self.database = database self.table = table self.kcsbData = KustoConnectionStringBuilder.with_aad_application_key_authentication(f"https://{cluster}:443/", clientId, clientSecret, authority) self.kcsbIngest = KustoConnectionStringBuilder.with_aad_application_key_authentication(f"https://ingest-{cluster}:443/", clientId, clientSecret, authority) self.dataClient = KustoClient(self.kcsbData) self.ingestClient = QueuedIngestClient(self.kcsbIngest) self.ingestionProps = IngestionProperties(database=database, table=table,)
def get_dataset(): in_use = microsoft_config kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication( in_use['cluster'], in_use['client_id'], in_use["client_secret"], in_use["authority_id"]) query = ''' RawEventsAzCli | where EventTimestamp > ago(2h) and EventTimestamp < ago(1h) // Remove pre 2.0.28 telemetry as it has a different schema | where toint(split(tostring(parse_json(Properties).["context.default.vs.core.telemetryapi.productversion"]), '.')[2]) > 28 | where tostring(parse_json(Properties).["reserved.datamodel.action.result"]) == "Success" | sort by UserId | project UserId,command=parse_json(Properties)["context.default.azurecli.rawcommand"], params = parse_json(Properties).["context.default.azurecli.params"], os_type = tostring(parse_json(Properties).["context.default.vs.core.os.type"]), source = tostring(parse_json(Properties).["context.default.azurecli.source"]), start_time_str= todatetime(parse_json(Properties).["context.default.azurecli.starttime"]), end_time_str = todatetime(parse_json(Properties).["context.default.azurecli.endtime"]) | take 500 ''' client = KustoClient(kcsb) response = client.execute("AzureCli", query) dataset = {} for row in response.primary_results[0]: uid, command, params, os_type, source, begin, end = row params = sort_params(params) if uid not in dataset: dataset[uid] = [(command, params, os_type, source, begin, end)] else: last_cmd, last_params, *_ = dataset[uid][-1] last_params = sort_params(last_params) if last_cmd == command and last_params == params: continue # duplicate dataset[uid].append( (command, params, os_type, source, begin, end)) return dataset
def setup_class(cls): # DM CS can be composed from engine CS cls.engine_cs = os.environ.get("ENGINE_CONNECTION_STRING") cls.dm_cs = os.environ.get( "DM_CONNECTION_STRING") or cls.engine_cs.replace( "//", "//ingest-") cls.app_id = os.environ.get("APP_ID") cls.app_key = os.environ.get("APP_KEY") cls.auth_id = os.environ.get("AUTH_ID") cls.test_db = os.environ.get("TEST_DATABASE") if not all([ cls.engine_cs, cls.dm_cs, cls.app_id, cls.app_key, cls.auth_id, cls.test_db ]): raise unittest.SkipTest("E2E environment is missing") # Init clients python_version = "_".join([str(v) for v in sys.version_info[:3]]) cls.test_table = "python_test_{0}_{1}_{2}".format( python_version, str(int(time.time())), random.randint(1, 100000)) cls.client = KustoClient(cls.engine_kcsb_from_env()) cls.ingest_client = KustoIngestClient(cls.dm_kcsb_from_env()) cls.streaming_ingest_client = KustoStreamingIngestClient( cls.engine_kcsb_from_env()) cls.input_folder_path = cls.get_file_path() cls.csv_file_path = os.path.join(cls.input_folder_path, "dataset.csv") cls.tsv_file_path = os.path.join(cls.input_folder_path, "dataset.tsv") cls.zipped_csv_file_path = os.path.join(cls.input_folder_path, "dataset.csv.gz") cls.json_file_path = os.path.join(cls.input_folder_path, "dataset.json") cls.zipped_json_file_path = os.path.join(cls.input_folder_path, "dataset.jsonz.gz") cls.current_count = 0 cls.client.execute( cls.test_db, ".create table {0} (rownumber: int, rowguid: string, xdouble: real, xfloat: real, xbool: bool, xint16: int, xint32: int, xint64: long, xuint8: long, xuint16: long, xuint32: long, xuint64: long, xdate: datetime, xsmalltext: string, xtext: string, xnumberAsText: string, xtime: timespan, xtextWithNulls: string, xdynamicWithNulls: dynamic)" .format(cls.test_table), ) cls.client.execute( cls.test_db, ".create table {0} ingestion json mapping 'JsonMapping' {1}". format(cls.test_table, cls.test_table_json_mapping_reference()))
def _set_client(self): self.client = KustoClient( kusto_cluster=self.cluster_url_template.format( self._parsed_conn.get('cluster')), client_id=self._parsed_conn.get('clientid'), client_secret=self._parsed_conn.get('clientsecret'), username=self._parsed_conn.get('username'), password=self._parsed_conn.get('password'), authority=self._parsed_conn.get('tenant')) # patch that replace the authetication helper my_aad_helper = _MyAadHelper( kusto_cluster=self.cluster_url_template.format( self._parsed_conn.get('cluster')), client_id=self._parsed_conn.get('clientid'), client_secret=self._parsed_conn.get('clientsecret'), username=self._parsed_conn.get('username'), password=self._parsed_conn.get('password'), authority=self._parsed_conn.get('tenant')) self.client._aad_helper = my_aad_helper
def get_kusto_client(server: str, database: str) -> KustoClient: """ Helper to get an authenticated KustoClient. Try to use Az CLI cached credentials, fall back to device code auth. :param server: The (short) name of a Kusto server cluster, not the full URI :param database: The name of the initial catalog to connect to """ logger = logging.getLogger(__name__) server_uri = f"https://{server}.kusto.windows.net" try: kcsb = KustoConnectionStringBuilder.with_az_cli_authentication( server_uri) client = KustoClient(kcsb) # hit the server to force authentication client.execute_query(database, "print('hi')") return client except KustoAuthenticationError: kcsb = KustoConnectionStringBuilder.with_aad_device_authentication( server_uri) client = KustoClient(kcsb) client.execute_query(database, "print('hi')") return client