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"
)
