def ingest_to_ADX(filepath, filesize): KCSB_INGEST = KustoConnectionStringBuilder.with_aad_device_authentication( DATA_INGESTION_URI) KCSB_INGEST.authority_id = AAD_TENANT_ID KCSB_ENGINE = KustoConnectionStringBuilder.with_aad_device_authentication( URI) KCSB_ENGINE.authority_id = AAD_TENANT_ID INGESTION_CLIENT = KustoIngestClient(KCSB_INGEST) INGESTION_PROPERTIES = IngestionProperties( database=DATABASE, table=DESTINATION_TABLE, dataFormat=DataFormat.CSV, mappingReference=DESTINATION_TABLE_COLUMN_MAPPING, additionalProperties={'ignoreFirstRecord': 'true'}, reportLevel=ReportLevel.FailuresAndSuccesses) BLOB_PATH = "https://" + SOURCE_CSV_BLOB_ACCOUNT + ".blob.core.windows.net/" + SOURCE_CSV_CONTAINER + "/" + filepath + SOURCE_CSV_BLOB_TOKEN BLOB_DESCRIPTOR = BlobDescriptor( BLOB_PATH, filesize) # 10 is the raw size of the data in bytes INGESTION_CLIENT.ingest_from_blob( BLOB_DESCRIPTOR, ingestion_properties=INGESTION_PROPERTIES) print('Done queuing up ingestion with Azure Data Explorer ' + filepath)
def __init__( self, kusto_cluster, client_id=None, client_secret=None, username=None, password=None, certificate=None, certificate_thumbprint=None, authority=None, ): """ Kusto Client constructor. Parameters ---------- kusto_cluster : str Kusto cluster endpoint. Example: https://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 authority : 'microsoft.com', optional In case your tenant is not microsoft please use this param. """ if all([username, password]): kcsb = KustoConnectionStringBuilder.with_aad_user_password_authentication( kusto_cluster, username, password) elif all([client_id, client_secret]): kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication( kusto_cluster, client_id, client_secret) elif all([client_id, certificate, certificate_thumbprint]): kcsb = KustoConnectionStringBuilder.with_aad_application_certificate_authentication( kusto_cluster, client_id, certificate, certificate_thumbprint) else: kcsb = KustoConnectionStringBuilder.with_aad_device_authentication( kusto_cluster) if authority: kcsb.authority_id = authority self.client = KustoClient(kcsb) # replace aadhelper to use remote browser in interactive mode self.client._aad_helper = _MyAadHelper(kcsb) self.mgmt_endpoint_version = "v2" if self.client._mgmt_endpoint.endswith( "v2/rest/query") else "v1" self.query_endpoint_version = "v2" if self.client._query_endpoint.endswith( "v2/rest/query") else "v1"
def Ingest(Tag): # setting AUTHORITY_ID = "6babcaad-604b-40ac-a9d7-9fd97c0b779f" INGESTCLUSTER = "https://ingest-cgadataout.kusto.windows.net" KUSTOCLUSTER = "https://cgadataout.kusto.windows.net" DATABASE = "DevRelWorkArea" # Create table KCSB_DATA = KustoConnectionStringBuilder.with_aad_device_authentication( KUSTOCLUSTER) DESTINATION_TABLE = "RepoContributors" DESTINATION_TABLE_COLUMN_MAPPING = "RepoContributors_CSV_Mapping" KUSTO_CLIENT = KustoClient(KCSB_DATA) DROP_TABLE_IF_EXIST = ".drop table RepoContributors ifexists" RESPONSE = KUSTO_CLIENT.execute_mgmt(DATABASE, DROP_TABLE_IF_EXIST) CREATE_TABLE_COMMAND = ".create table RepoContributors (Article: string, Contributors: int64, Data: string)" RESPONSE = KUSTO_CLIENT.execute_mgmt(DATABASE, CREATE_TABLE_COMMAND) print("RepoContributors table is created") # Create mapping CREATE_MAPPING_COMMAND = """.create table RepoContributors ingestion csv mapping 'RepoContributors_CSV_Mapping' '[{"Name": "Article","datatype": "string","Ordinal": 0},{"Name": "Contributors","datatype": "int64","Ordinal": 1},{"Name": "Data","datatype": "string","Ordinal": 2}]'""" RESPONSE = KUSTO_CLIENT.execute_mgmt(DATABASE, CREATE_MAPPING_COMMAND) print("mapping is created") # Ingest # The authentication method will be taken from the chosen KustoConnectionStringBuilder. ingestion_props = IngestionProperties( database="DevRelWorkArea", table="RepoContributors", dataFormat=DataFormat.CSV, ingestByTags=[Tag], dropByTags=[Tag], mappingReference=DESTINATION_TABLE_COLUMN_MAPPING, reportLevel=ReportLevel.FailuresAndSuccesses, additionalProperties={'ignoreFirstRecord': 'true'}) kcsb = KustoConnectionStringBuilder.with_aad_device_authentication( INGESTCLUSTER) client = KustoIngestClient(kcsb) # ingest from file file_descriptor = FileDescriptor( r"D:\test\Results\log_data_merge\merge_microsoftdocs_sql-docs-pr.txt", 3333) # 3333 is the raw size of the data in bytes. client.ingest_from_file(file_descriptor, ingestion_properties=ingestion_props) # if status updates are required, something like this can be done return 1
def get_conn(self) -> KustoClient: """Return a KustoClient object.""" conn = self.get_connection(self.conn_id) cluster = conn.host if not cluster: raise AirflowException('Host connection option is required') def get_required_param(name: str) -> str: """Extract required parameter from extra JSON, raise exception if not found""" value = conn.extra_dejson.get(name) if not value: raise AirflowException( f'Extra connection option is missing required parameter: `{name}`' ) return value auth_method = get_required_param('auth_method') or get_required_param( 'extra__azure_data_explorer__auth_method') if auth_method == 'AAD_APP': tenant = get_required_param('tenant') or get_required_param( 'extra__azure_data_explorer__tenant') kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication( cluster, conn.login, conn.password, tenant) elif auth_method == 'AAD_APP_CERT': certificate = get_required_param( 'certificate') or get_required_param( 'extra__azure_data_explorer__certificate') thumbprint = get_required_param( 'thumbprint') or get_required_param( 'extra__azure_data_explorer__thumbprint') tenant = get_required_param('tenant') or get_required_param( 'extra__azure_data_explorer__tenant') kcsb = KustoConnectionStringBuilder.with_aad_application_certificate_authentication( cluster, conn.login, certificate, thumbprint, tenant, ) elif auth_method == 'AAD_CREDS': tenant = get_required_param('tenant') or get_required_param( 'extra__azure_data_explorer__tenant') kcsb = KustoConnectionStringBuilder.with_aad_user_password_authentication( cluster, conn.login, conn.password, tenant) elif auth_method == 'AAD_DEVICE': kcsb = KustoConnectionStringBuilder.with_aad_device_authentication( cluster) else: raise AirflowException( f'Unknown authentication method: {auth_method}') return KustoClient(kcsb)
def test_no_credentials(self): """Checks kcsb that is created with no credentials""" kcsbs = [ KustoConnectionStringBuilder("localhost"), KustoConnectionStringBuilder("data Source=localhost"), KustoConnectionStringBuilder("Addr=localhost"), KustoConnectionStringBuilder("Addr = localhost"), KustoConnectionStringBuilder.with_aad_device_authentication( "localhost"), ] for kcsb in kcsbs: self._validate_kcsb_without_credentials(kcsb, "localhost")
def test_add_token_provider(self): caller_token = "caller token" token_provider = lambda: caller_token kscb = KustoConnectionStringBuilder.with_token_provider("localhost", token_provider) assert kscb.token_provider() == caller_token exception_occurred = False try: kscb = KustoConnectionStringBuilder.with_token_provider("localhost", caller_token) except AssertionError as ex: exception_occurred = True finally: assert exception_occurred
def ingest_to_ADX(filepath, telemetry_block_blob_service, container_name, blob_account, file_size, tc,vm_uuid,deploy_uuid,config_uuid): ingest_source_id=str(uuid.uuid4()) KCSB_INGEST = KustoConnectionStringBuilder.with_aad_device_authentication(DATA_INGESTION_URI) KCSB_INGEST.authority_id = APP_AAD_TENANT_ID INGESTION_CLIENT = KustoIngestClient(KCSB_INGEST) ing_map=[JsonColumnMapping("vm_uuid", "$.vm_uuid", "string"), JsonColumnMapping("deploy_uuid", "$.deployment_description[0].deploy_uuid", "string"), JsonColumnMapping("config_uuid", "$.vm_configuration[0].config_uuid", "string"), JsonColumnMapping("rawdata", "$", "dynamic")] INGESTION_PROPERTIES = IngestionProperties(database=DATABASE, table=DESTINATION_TABLE, dataFormat=DataFormat.JSON, ingestionMapping=ing_map, reportLevel=ReportLevel.FailuresAndSuccesses,flushImmediately=IS_FLUSH_IMMEDIATELY) print("Database {} Tabele {}".format(DATABASE,DESTINATION_TABLE)) BLOB_PATH = "https://" + blob_account + ".blob.core.windows.net/" + container_name + "/" + filepath + CLEAN_FILE_TOKEN print (BLOB_PATH,' ',str(file_size), ingest_source_id) BLOB_DESCRIPTOR = BlobDescriptor(BLOB_PATH, file_size, ingest_source_id) # 10 is the raw size of the data in bytes INGESTION_CLIENT.ingest_from_blob(BLOB_DESCRIPTOR,ingestion_properties=INGESTION_PROPERTIES) tc.context.properties["ingest_source_id"]=ingest_source_id min_datatime=0 max_datatime=0 total_records=1 doc_id=save_COSMOS_log(vm_uuid,deploy_uuid,config_uuid,filepath,min_datatime,max_datatime, total_records,ingest_source_id,blob_account,container_name, tc) tc.track_event(APP_INSIGHT_INGEST_EVENT_NAME, { 'FILE_PATH': filepath,'DOC_ID':doc_id,"SOURCE_ID":ingest_source_id }, { 'TOTOAL_RECORDS': total_records, 'FILE_SIZE':file_size,'MIN_DATETIME':min_datatime,'MAX_DATETIME': max_datatime }) log_msg="{} Done queuing up ingestion with Azure Data Explorer {}, Ingest SourceID {}".format(LOG_MESSAGE_HEADER,filepath,ingest_source_id) print(log_msg) tc.track_trace(log_msg) tc.flush()
def test_aad_user_with_authority(self): """Checks kcsb that is created with AAD user credentials.""" user = "******" password = "******" authority_id = "13456" kcsb = KustoConnectionStringBuilder.with_aad_user_password_authentication( "localhost", user, password, authority_id) self.assertEqual(kcsb.data_source, "localhost") self.assertTrue(kcsb.aad_federated_security) self.assertEqual(kcsb.aad_user_id, user) self.assertEqual(kcsb.password, password) self.assertIsNone(kcsb.application_client_id) self.assertIsNone(kcsb.application_key) self.assertEqual(kcsb.authority_id, authority_id) self.assertEqual( repr(kcsb), "Data Source=localhost;AAD Federated Security=True;AAD User ID={0};Password={1};Authority Id=13456" .format(user, password), ) self.assertEqual( str(kcsb), "Data Source=localhost;AAD Federated Security=True;AAD User ID={0};Password={1};Authority Id=13456" .format(user, self.PASSWORDS_REPLACEMENT), )
def createKustoConnection(uri, tenantID, appID, appKey): connectionString = None try: connectionString = KustoConnectionStringBuilder.with_aad_application_key_authentication(uri, appID, appKey, tenantID) except Exception as e: logging.error("Could not create a connection string:%s"%e) return connectionString
def test_add_msi(self): client_guid = "kjhjk" object_guid = "87687687" res_guid = "kajsdghdijewhag" kcsb = [ KustoConnectionStringBuilder.with_aad_managed_service_identity_authentication("localhost0", timeout=1), KustoConnectionStringBuilder.with_aad_managed_service_identity_authentication("localhost1", client_id=client_guid, timeout=2), KustoConnectionStringBuilder.with_aad_managed_service_identity_authentication("localhost2", object_id=object_guid, timeout=3), KustoConnectionStringBuilder.with_aad_managed_service_identity_authentication("localhost3", msi_res_id=res_guid), ] assert kcsb[0].msi_authentication assert kcsb[0].msi_parameters["resource"] == "localhost0" assert kcsb[0].msi_parameters["timeout"] == 1 assert "client_id" not in kcsb[0].msi_parameters assert "object_id" not in kcsb[0].msi_parameters assert "msi_res_id" not in kcsb[0].msi_parameters assert kcsb[1].msi_authentication assert kcsb[1].msi_parameters["resource"] == "localhost1" assert kcsb[1].msi_parameters["timeout"] == 2 assert kcsb[1].msi_parameters["client_id"] == client_guid assert "object_id" not in kcsb[1].msi_parameters assert "msi_res_id" not in kcsb[1].msi_parameters assert kcsb[2].msi_authentication assert kcsb[2].msi_parameters["resource"] == "localhost2" assert kcsb[2].msi_parameters["timeout"] == 3 assert "client_id" not in kcsb[2].msi_parameters assert kcsb[2].msi_parameters["object_id"] == object_guid assert "msi_res_id" not in kcsb[2].msi_parameters assert kcsb[3].msi_authentication assert kcsb[3].msi_parameters["resource"] == "localhost3" assert "timeout" not in kcsb[3].msi_parameters assert "client_id" not in kcsb[3].msi_parameters assert "object_id" not in kcsb[3].msi_parameters assert kcsb[3].msi_parameters["msi_res_id"] == res_guid exception_occurred = False try: fault = KustoConnectionStringBuilder.with_aad_managed_service_identity_authentication("localhost", client_id=client_guid, object_id=object_guid) except ValueError as e: exception_occurred = True finally: assert exception_occurred
def ingest_to_ADX(filepath, telemetry_block_blob_service, container_name, blob_account, tc): ingest_source_id = str(uuid.uuid4()) #file_size=BlockBlobService.get_blob_properties(telemetry_block_blob_service,container_name,filepath).properties.content_length #print (filepath+" File Size "+str(file_size)) KCSB_INGEST = KustoConnectionStringBuilder.with_aad_device_authentication( DATA_INGESTION_URI) KCSB_INGEST.authority_id = APP_AAD_TENANT_ID vm_uuid, config_uuid, deploy_uuid, file_size, min_datatime, max_datatime, total_records = get_uuids_from_csv( telemetry_block_blob_service, container_name, filepath) dropByTag = vm_uuid + '_' + config_uuid + '_' + deploy_uuid INGESTION_CLIENT = KustoIngestClient(KCSB_INGEST) INGESTION_PROPERTIES = IngestionProperties( database=DATABASE, table=DESTINATION_TABLE, dataFormat=DataFormat.CSV, mappingReference=DESTINATION_TABLE_COLUMN_MAPPING, additionalProperties={ 'ignoreFirstRecord': 'true', 'reportMethod': 'QueueAndTable' }, reportLevel=ReportLevel.FailuresAndSuccesses, dropByTags=[dropByTag], flushImmediately=IS_FLUSH_IMMEDIATELY) BLOB_PATH = "https://" + SOURCE_OSMETRICS_BLOB_ACCOUNT + ".blob.core.windows.net/" + SOURCE_OSMETRICS_CONTAINER + "/" + filepath + SOURCE_OSMETRICS_FILE_TOKEN #print (BLOB_PATH,' ',str(file_size)) BLOB_DESCRIPTOR = BlobDescriptor( BLOB_PATH, file_size, ingest_source_id) # 10 is the raw size of the data in bytes INGESTION_CLIENT.ingest_from_blob( BLOB_DESCRIPTOR, ingestion_properties=INGESTION_PROPERTIES) tc.context.properties["ingest_source_id"] = str(ingest_source_id) doc_id = save_COSMOS_log(vm_uuid, deploy_uuid, config_uuid, filepath, min_datatime, max_datatime, total_records, ingest_source_id, blob_account, container_name, tc) tc.track_event(APP_INSIGHT_INGEST_EVENT_NAME, { 'FILE_PATH': filepath, 'DOC_ID': doc_id, "SOURCE_ID": ingest_source_id }, { 'TOTOAL_RECORDS': total_records, 'FILE_SIZE': file_size, 'MIN_DATETIME': min_datatime, 'MAX_DATETIME': max_datatime }) log_msg = "{} Done queuing up ingestion with Azure Data Explorer {}, Ingest SourceID {}".format( LOG_MESSAGE_HEADER, filepath, ingest_source_id) print(log_msg) tc.track_trace(log_msg) tc.flush()
def test_aad_user(self): """Checks kcsb that is created with AAD user credentials.""" user = "******" password = "******" kcsbs = [ KustoConnectionStringBuilder( "localhost;AAD User ID={0};password={1}".format( user, password)), KustoConnectionStringBuilder( "Data Source=localhost ; AaD User ID={0}; Password ={1}". format(user, password)), KustoConnectionStringBuilder( " Addr = localhost ; AAD User ID = {0} ; Pwd ={1}".format( user, password)), KustoConnectionStringBuilder( "Network Address = localhost; AAD User iD = {0} ; Pwd = {1} ". format(user, password)), KustoConnectionStringBuilder.with_aad_user_password_authentication( "localhost", user, password), ] kcsb1 = KustoConnectionStringBuilder("Server=localhost") kcsb1[KustoConnectionStringBuilder.ValidKeywords.aad_user_id] = user kcsb1[KustoConnectionStringBuilder.ValidKeywords.password] = password kcsbs.append(kcsb1) kcsb2 = KustoConnectionStringBuilder("server=localhost") kcsb2["AAD User ID"] = user kcsb2["Password"] = password kcsbs.append(kcsb2) for kcsb in kcsbs: self._validate_kcsb_with_aad_user(kcsb, "localhost", user, password)
def __init__(self, conn_kv): """ Kusto Client constructor. Parameters ---------- kusto_cluster : str Kusto cluster endpoint. Example: https://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 authority : 'microsoft.com', optional In case your tenant is not microsoft please use this param. """ kusto_cluster = "https://{0}.kusto.windows.net".format(conn_kv["cluster"]) if all([conn_kv.get("username"), conn_kv.get("password")]): kcsb = KustoConnectionStringBuilder.with_aad_user_password_authentication(kusto_cluster, conn_kv.get("username"), conn_kv.get("password")) if conn_kv.get("tenant") is not None: kcsb.authority_id = conn_kv.get("tenant") elif all([conn_kv.get("clientid"), conn_kv.get("clientsecret")]): kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication( kusto_cluster, conn_kv.get("clientid"), conn_kv.get("clientsecret"), conn_kv.get("tenant")) elif all([conn_kv.get("clientid"), conn_kv.get("certificate"), conn_kv.get("certificate_thumbprint")]): kcsb = KustoConnectionStringBuilder.with_aad_application_certificate_authentication( kusto_cluster, conn_kv.get("clientid"), conn_kv.get("certificate"), conn_kv.get("certificate_thumbprint", conn_kv.get("tenant")) ) else: kcsb = KustoConnectionStringBuilder.with_aad_device_authentication(kusto_cluster) if conn_kv.get("tenant") is not None: kcsb.authority_id = conn_kv.get("tenant") self.client = KustoClient(kcsb) # replace aadhelper to use remote browser in interactive mode self.client._aad_helper = _MyAadHelper(kcsb, self._DEFAULT_CLIENTID) self.mgmt_endpoint_version = "v2" if self.client._mgmt_endpoint.endswith("v2/rest/query") else "v1" self.query_endpoint_version = "v2" if self.client._query_endpoint.endswith("v2/rest/query") else "v1"
def test_token_provider_auth(): valid_token_provider = lambda: "caller token" invalid_token_provider = lambda: 12345678 valid_kcsb = KustoConnectionStringBuilder.with_token_provider("localhost", valid_token_provider) invalid_kcsb = KustoConnectionStringBuilder.with_token_provider("localhost", invalid_token_provider) valid_helper = _AadHelper(valid_kcsb) invalid_helper = _AadHelper(invalid_kcsb) auth_header = valid_helper.acquire_authorization_header() assert auth_header.index(valid_token_provider()) > -1 try: invalid_helper.acquire_authorization_header() except KustoAuthenticationError as e: assert e.authentication_method == AuthenticationMethod.token_provider.value assert str(e.exception).index(str(type(invalid_token_provider()))) > -1
def get_kusto_client() -> KustoIngestClient: cluster = "https://ingest-" + os.environ['KUSTO_CLUSTER'] + ".kusto.windows.net" username = os.environ['KUSTO_USERNAME'] password = os.environ['KUSTO_PASSWORD'] authority_id = os.environ['KUSTO_TENANT_ID'] kcsb = KustoConnectionStringBuilder.with_aad_user_password_authentication(cluster, username, password, authority_id) return KustoIngestClient(kcsb)
def test_no_credentials(self): """Checks kcsb that is created with no credentials""" kcsbs = [ KustoConnectionStringBuilder("localhost"), KustoConnectionStringBuilder("data Source=localhost"), KustoConnectionStringBuilder("Addr=localhost"), KustoConnectionStringBuilder("Addr = localhost"), ] for kcsb in kcsbs: assert kcsb.data_source == "localhost" assert not kcsb.aad_federated_security assert kcsb.aad_user_id is None assert kcsb.password is None assert kcsb.application_client_id is None assert kcsb.application_key is None assert kcsb.authority_id == "common" assert repr(kcsb) == "Data Source=localhost;Authority Id=common" assert str(kcsb) == "Data Source=localhost;Authority Id=common"
def get_conn(self) -> KustoClient: """Return a KustoClient object.""" conn = self.get_connection(self.conn_id) cluster = conn.host if not cluster: raise AirflowException('Host connection option is required') def get_required_param(name): """Extract required parameter from extra JSON, raise exception if not found""" value = conn.extra_dejson.get(name) if not value: raise AirflowException( 'Extra connection option is missing required parameter: `{}`' .format(name)) return value auth_method = get_required_param('auth_method') if auth_method == 'AAD_APP': kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication( cluster, conn.login, conn.password, get_required_param('tenant')) elif auth_method == 'AAD_APP_CERT': kcsb = KustoConnectionStringBuilder.with_aad_application_certificate_authentication( cluster, conn.login, get_required_param('certificate'), get_required_param('thumbprint'), get_required_param('tenant'), ) elif auth_method == 'AAD_CREDS': kcsb = KustoConnectionStringBuilder.with_aad_user_password_authentication( cluster, conn.login, conn.password, get_required_param('tenant')) elif auth_method == 'AAD_DEVICE': kcsb = KustoConnectionStringBuilder.with_aad_device_authentication( cluster) else: raise AirflowException( 'Unknown authentication method: {}'.format(auth_method)) return KustoClient(kcsb)
def test_conn_method_aad_device(self, mock_init): mock_init.return_value = None db.merge_conn( Connection(conn_id=ADX_TEST_CONN_ID, conn_type='azure_data_explorer', host='https://help.kusto.windows.net', extra=json.dumps({'auth_method': 'AAD_DEVICE'}))) AzureDataExplorerHook(azure_data_explorer_conn_id=ADX_TEST_CONN_ID) assert mock_init.called_with( KustoConnectionStringBuilder.with_aad_device_authentication( 'https://help.kusto.windows.net'))
def __init__(self, trainingConfig, trainingId): self.trainingId = trainingId cluster = "https://usage360.kusto.windows.net" authority_id = "72f988bf-86f1-41af-91ab-2d7cd011db47" client_id = kustoClientId client_secret = kustoClientSecret kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication(cluster, client_id, client_secret, authority_id) self.kustoClient = KustoClient(kcsb) self.garbageList = [x.strip() for x in open("metadata/garbagePhrases.txt", "r").readlines()] self.striptrailers = [x.strip() for x in open("metadata/stripTrailers.txt", "r").readlines()] self.shortPhrases = [x.strip() for x in open("metadata/shortPhrasesList.txt", "r").readlines()] self.trainingConfig = trainingConfig
def test_no_credentials(self): """Checks kcsb that is created with no credentials""" kcsbs = [ KustoConnectionStringBuilder("localhost"), KustoConnectionStringBuilder("data Source=localhost"), KustoConnectionStringBuilder("Addr=localhost"), KustoConnectionStringBuilder("Addr = localhost"), ] for kcsb in kcsbs: self.assertEqual(kcsb.data_source, "localhost") self.assertFalse(kcsb.aad_federated_security) self.assertIsNone(kcsb.aad_user_id) self.assertIsNone(kcsb.password) self.assertIsNone(kcsb.application_client_id) self.assertIsNone(kcsb.application_key) self.assertEqual(kcsb.authority_id, "common") self.assertEqual(repr(kcsb), "Data Source=localhost;Authority Id=common") self.assertEqual(str(kcsb), "Data Source=localhost;Authority Id=common")
def test_aad_device_login(self): """Checks kcsb that is created with AAD device login.""" kcsb = KustoConnectionStringBuilder.with_aad_device_authentication("localhost") assert kcsb.data_source == "localhost" assert kcsb.aad_federated_security assert kcsb.aad_user_id is None assert kcsb.password is None assert kcsb.application_client_id is None assert kcsb.application_key is None assert kcsb.authority_id == "common" assert repr(kcsb) == "Data Source=localhost;AAD Federated Security=True;Authority Id=common" assert str(kcsb) == "Data Source=localhost;AAD Federated Security=True;Authority Id=common"
def test_aad_app(self): """Checks kcsb that is created with AAD application credentials.""" uuid = str(uuid4()) key = "key of application" kcsbs = [ KustoConnectionStringBuilder( "localhost;Application client Id={0};application Key={1}". format(uuid, key)), KustoConnectionStringBuilder( "Data Source=localhost ; Application Client Id={0}; Appkey ={1}" .format(uuid, key)), KustoConnectionStringBuilder( " Addr = localhost ; AppClientId = {0} ; AppKey ={1}".format( uuid, key)), KustoConnectionStringBuilder( "Network Address = localhost; AppClientId = {0} ; AppKey ={1}". format(uuid, key)), KustoConnectionStringBuilder. with_aad_application_key_authentication("localhost", uuid, key), ] kcsb1 = KustoConnectionStringBuilder("server=localhost") kcsb1[KustoConnectionStringBuilder.ValidKeywords. application_client_id] = uuid kcsb1[KustoConnectionStringBuilder.ValidKeywords.application_key] = key kcsbs.append(kcsb1) kcsb2 = KustoConnectionStringBuilder("Server=localhost") kcsb2["AppclientId"] = uuid kcsb2["Application key"] = key kcsbs.append(kcsb2) for kcsb in kcsbs: self._validate_kcsb_with_aad_app(kcsb, "localhost", uuid, key)
def get_client(cluster): """ get cached, authenticated client for given cluster """ global _client_cache c = _client_cache.get(cluster) if c is None: c = KustoClient( KustoConnectionStringBuilder.with_aad_device_authentication( cluster)) c.execute('VSO', 'print "a" | take 0') _client_cache[cluster] = c return c
def test_unauthorized_exception(): """Test the exception thrown when authorization fails.""" cluster = "https://somecluster.kusto.windows.net" username = "******" kcsb = KustoConnectionStringBuilder.with_aad_user_password_authentication(cluster, username, "StrongestPasswordEver", "authorityName") aad_helper = _AadHelper(kcsb) try: aad_helper.acquire_authorization_header() except KustoAuthenticationError as error: assert error.authentication_method == AuthenticationMethod.aad_username_password.value assert error.authority == "https://login.microsoftonline.com/authorityName" assert error.kusto_cluster == cluster assert error.kwargs["username"] == username
def test_aad_user(self): """Checks kcsb that is created with AAD user credentials.""" user = "******" password = "******" kcsbs = [ KustoConnectionStringBuilder( "localhost;AAD User ID={0};password={1} ;AAD Federated Security=True " .format(user, password)), KustoConnectionStringBuilder( "Data Source=localhost ; AaD User ID={0}; Password ={1} ;AAD Federated Security=True" .format(user, password)), KustoConnectionStringBuilder( " Addr = localhost ; AAD User ID = {0} ; Pwd ={1} ;AAD Federated Security=True" .format(user, password)), KustoConnectionStringBuilder( "Network Address = localhost; AAD User iD = {0} ; Pwd = {1} ;AAD Federated Security= True " .format(user, password)), KustoConnectionStringBuilder.with_aad_user_password_authentication( "localhost", user, password), ] kcsb1 = KustoConnectionStringBuilder("Server=localhost") kcsb1[KustoConnectionStringBuilder.ValidKeywords.aad_user_id] = user kcsb1[KustoConnectionStringBuilder.ValidKeywords.password] = password kcsb1[KustoConnectionStringBuilder.ValidKeywords. aad_federated_security] = True kcsbs.append(kcsb1) kcsb2 = KustoConnectionStringBuilder("server=localhost") kcsb2["AAD User ID"] = user kcsb2["Password"] = password kcsb2["aad federated security"] = True kcsbs.append(kcsb2) for kcsb in kcsbs: self.assertEqual(kcsb.data_source, "localhost") self.assertTrue(kcsb.aad_federated_security) self.assertEqual(kcsb.aad_user_id, user) self.assertEqual(kcsb.password, password) self.assertIsNone(kcsb.application_client_id) self.assertIsNone(kcsb.application_key) self.assertEqual(kcsb.authority_id, "common") self.assertEqual( repr(kcsb), "Data Source=localhost;AAD Federated Security=True;AAD User ID={0};Password={1};Authority Id=common" .format(user, password), ) self.assertEqual( str(kcsb), "Data Source=localhost;AAD Federated Security=True;AAD User ID={0};Password={1};Authority Id=common" .format(user, self.PASSWORDS_REPLACEMENT), )
def test_msi_auth(): """ * * * Note * * * Each connection test takes about 15-20 seconds which is the time it takes TCP to fail connecting to the nonexistent MSI endpoint The timeout option does not seem to affect this behavior. Could be it only affects the waiting time fora response in successful connections. Please be prudent in adding any future tests! """ client_guid = "kjhjk" object_guid = "87687687" res_guid = "kajsdghdijewhag" kcsb = [ KustoConnectionStringBuilder.with_aad_managed_service_identity_authentication("localhost", timeout=1), KustoConnectionStringBuilder.with_aad_managed_service_identity_authentication("localhost", client_id=client_guid, timeout=1), KustoConnectionStringBuilder.with_aad_managed_service_identity_authentication("localhost", object_id=object_guid, timeout=1), KustoConnectionStringBuilder.with_aad_managed_service_identity_authentication("localhost", msi_res_id=res_guid, timeout=1), ] helpers = [_AadHelper(kcsb[0]), _AadHelper(kcsb[1]), _AadHelper(kcsb[2]), _AadHelper(kcsb[3])] try: helpers[0].acquire_authorization_header() except KustoAuthenticationError as e: assert e.authentication_method == AuthenticationMethod.aad_msi.value assert "client_id" not in e.kwargs assert "object_id" not in e.kwargs assert "msi_res_id" not in e.kwargs try: helpers[1].acquire_authorization_header() except KustoAuthenticationError as e: assert e.authentication_method == AuthenticationMethod.aad_msi.value assert e.kwargs["client_id"] == client_guid assert "object_id" not in e.kwargs assert "msi_res_id" not in e.kwargs assert str(e.exception).index("client_id") > -1 assert str(e.exception).index(client_guid) > -1
def test_aad_user_token(self): """Checks kcsb that is created with AAD user token.""" token = "The user hardest token ever" kcsb = KustoConnectionStringBuilder.with_aad_user_token_authentication("localhost", user_token=token) assert kcsb.data_source == "localhost" assert kcsb.user_token == token assert kcsb.aad_federated_security assert kcsb.aad_user_id is None assert kcsb.password is None assert kcsb.application_client_id is None assert kcsb.application_key is None assert kcsb.application_token is None assert kcsb.authority_id == "common" assert repr(kcsb) == "Data Source=localhost;AAD Federated Security=True;Authority Id=common;User Token=%s" % token assert str(kcsb) == "Data Source=localhost;AAD Federated Security=True;Authority Id=common;User Token=%s" % self.PASSWORDS_REPLACEMENT
def run_query(self, query, user): kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication( connection_string=self.configuration["cluster"], aad_app_id=self.configuration["azure_ad_client_id"], app_key=self.configuration["azure_ad_client_secret"], authority_id=self.configuration["azure_ad_tenant_id"], ) client = KustoClient(kcsb) db = self.configuration["database"] try: response = client.execute(db, query) result_cols = response.primary_results[0].columns result_rows = response.primary_results[0].rows columns = [] rows = [] for c in result_cols: columns.append( { "name": c.column_name, "friendly_name": c.column_name, "type": TYPES_MAP.get(c.column_type, None), } ) # rows must be [{'column1': value, 'column2': value}] for row in result_rows: rows.append(row.to_dict()) error = None data = {"columns": columns, "rows": rows} json_data = json_dumps(data) except KustoServiceError as err: json_data = None try: error = err.args[1][0]["error"]["@message"] except (IndexError, KeyError): error = err.args[1] except KeyboardInterrupt: json_data = None error = "Query cancelled by user." return json_data, error
def test_conn_method_aad_creds(self, mock_init): mock_init.return_value = None db.merge_conn( Connection(conn_id=ADX_TEST_CONN_ID, conn_type='azure_data_explorer', login='******', password='******', host='https://help.kusto.windows.net', extra=json.dumps({ 'tenant': 'tenant', 'auth_method': 'AAD_CREDS' }))) AzureDataExplorerHook(azure_data_explorer_conn_id=ADX_TEST_CONN_ID) assert mock_init.called_with( KustoConnectionStringBuilder.with_aad_user_password_authentication( 'https://help.kusto.windows.net', 'client_id', 'client secret', 'tenant'))
def test_aad_device_login(self): """Checks kcsb that is created with AAD device login.""" kcsb = KustoConnectionStringBuilder.with_aad_device_authentication( "localhost") self.assertEqual(kcsb.data_source, "localhost") self.assertTrue(kcsb.aad_federated_security) self.assertIsNone(kcsb.aad_user_id) self.assertIsNone(kcsb.password) self.assertIsNone(kcsb.application_client_id) self.assertIsNone(kcsb.application_key) self.assertEqual(kcsb.authority_id, "common") self.assertEqual( repr(kcsb), "Data Source=localhost;AAD Federated Security=True;Authority Id=common" ) self.assertEqual( str(kcsb), "Data Source=localhost;AAD Federated Security=True;Authority Id=common" )