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 test_user_app_token_auth():
token = "123456446"
user_kcsb = KustoConnectionStringBuilder.with_aad_user_token_authentication(
KUSTO_TEST_URI, token)
app_kcsb = KustoConnectionStringBuilder.with_aad_application_token_authentication(
KUSTO_TEST_URI, token)
user_helper = _AadHelper(user_kcsb)
app_helper = _AadHelper(app_kcsb)
auth_header = user_helper.acquire_authorization_header()
assert auth_header.index(token) > -1
auth_header = app_helper.acquire_authorization_header()
assert auth_header.index(token) > -1
def authenticate_to_kusto_ingress(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 = KustoIngestClient(kcsb)
return kusto_client
def engine_kcsb_from_env() -> KustoConnectionStringBuilder:
engine_cs = os.environ.get("ENGINE_CONNECTION_STRING")
app_id = os.environ.get("APP_ID")
app_key = os.environ.get("APP_KEY")
auth_id = os.environ.get("AUTH_ID")
return KustoConnectionStringBuilder.with_aad_application_key_authentication(
engine_cs, app_id, app_key, auth_id)
def write_to_db(coverage_data, args):
# connect to database
cluster = "https://ingest-onnxruntimedashboarddb.southcentralus.kusto.windows.net"
kcsb = KustoConnectionStringBuilder.with_az_cli_authentication(cluster)
# The authentication method will be taken from the chosen KustoConnectionStringBuilder.
client = QueuedIngestClient(kcsb)
fields = [
"UploadTime", "CommitId", "Coverage", "LinesCovered", "TotalLines",
"OS", "Arch", "BuildConfig", "ReportURL", "Branch"
]
now_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
rows = [[
now_str, args.commit_hash, coverage_data['coverage'],
coverage_data['lines_covered'], coverage_data['lines_valid'],
args.os.lower(),
args.arch.lower(),
args.build_config.lower(),
args.report_url.lower(),
args.branch.lower()
]]
ingestion_props = IngestionProperties(
database="powerbi",
table="test_coverage",
data_format=DataFormat.CSV,
report_level=ReportLevel.FailuresAndSuccesses)
df = pandas.DataFrame(data=rows, columns=fields)
client.ingest_from_dataframe(df, ingestion_properties=ingestion_props)
def write_to_db(binary_size_data, args):
# connect to database
cluster = "https://ingest-onnxruntimedashboarddb.southcentralus.kusto.windows.net"
kcsb = KustoConnectionStringBuilder.with_az_cli_authentication(cluster)
# The authentication method will be taken from the chosen KustoConnectionStringBuilder.
client = QueuedIngestClient(kcsb)
fields = ["build_time", "build_id", "build_project", "commit_id", "os", "arch", "build_config", "size", "Branch"]
now_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
branch_name = os.environ.get("BUILD_SOURCEBRANCHNAME", "main")
rows = []
for row in binary_size_data:
rows.append(
[
now_str,
args.build_id,
args.build_project,
args.commit_hash,
row["os"],
row["arch"],
row["build_config"],
row["size"],
branch_name.lower(),
]
)
ingestion_props = IngestionProperties(
database="powerbi",
table="binary_size",
data_format=DataFormat.CSV,
report_level=ReportLevel.FailuresAndSuccesses,
)
df = pandas.DataFrame(data=rows, columns=fields)
client.ingest_from_dataframe(df, ingestion_properties=ingestion_props)
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 dm_kcsb_from_env() -> KustoConnectionStringBuilder:
engine_cs = os.environ.get("ENGINE_CONNECTION_STRING")
dm_cs = os.environ.get("DM_CONNECTION_STRING") or engine_cs.replace(
"//", "//ingest-")
app_id = os.environ.get("APP_ID")
app_key = os.environ.get("APP_KEY")
auth_id = os.environ.get("AUTH_ID")
return KustoConnectionStringBuilder.with_aad_application_key_authentication(
dm_cs, app_id, app_key, auth_id)
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 create_application_certificate_connection_string(
cls, cluster_url: str) -> KustoConnectionStringBuilder:
"""
Generates Kusto Connection String based on 'AppCertificate' Authentication Mode.
:param cluster_url: Url of cluster to connect to
:return: AppCertificate Kusto Connection String
"""
# TODO (config - optional): App ID & tenant, path to public certificate and path to private certificate pem file to authenticate with
app_id = os.environ.get("APP_ID")
app_tenant = os.environ.get("APP_TENANT")
private_key_pem_file_path = os.environ.get(
"PRIVATE_KEY_PEM_FILE_PATH")
cert_thumbprint = os.environ.get("CERT_THUMBPRINT")
public_cert_file_path = os.environ.get(
"PUBLIC_CERT_FILE_PATH"
) # Only used for "Subject Name and Issuer" auth
public_certificate = None
pem_certificate = None
try:
with open(private_key_pem_file_path, "r") as pem_file:
pem_certificate = pem_file.read()
except Exception as ex:
Utils.error_handler(
f"Failed to load PEM file from {private_key_pem_file_path}",
ex)
if public_cert_file_path:
try:
with open(public_cert_file_path, "r") as cert_file:
public_certificate = cert_file.read()
except Exception as ex:
Utils.error_handler(
f"Failed to load public certificate file from {public_cert_file_path}",
ex)
return KustoConnectionStringBuilder.with_aad_application_certificate_sni_authentication(
cluster_url, app_id, pem_certificate, public_certificate,
cert_thumbprint, app_tenant)
else:
return KustoConnectionStringBuilder.with_aad_application_certificate_authentication(
cluster_url, app_id, pem_certificate, cert_thumbprint,
app_tenant)
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 test_add_async_token_provider(self):
caller_token = "caller token"
async def async_token_provider():
return caller_token
kscb = KustoConnectionStringBuilder.with_async_token_provider(
"localhost", async_token_provider)
assert kscb.async_token_provider() is not None
exception_occurred = False
try:
kscb = KustoConnectionStringBuilder.with_async_token_provider(
"localhost", caller_token)
except AssertionError as ex:
exception_occurred = True
finally:
assert exception_occurred
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 __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 test_proxy_url_parsing(self):
"""Test Proxy URL Parsing"""
tests = {
"https://kusto.test.com":
"https://kusto.test.com",
"https://kusto.test.com/":
"https://kusto.test.com",
"https://kusto.test.com/test":
"https://kusto.test.com/test",
"https://kusto.test.com:4242":
"https://kusto.test.com:4242",
"https://kusto.test.com:4242/":
"https://kusto.test.com:4242",
"https://kusto.test.com:4242/test":
"https://kusto.test.com:4242/test",
"https://kusto.test.com;fed=true":
"https://kusto.test.com",
"https://kusto.test.com/;fed=true":
"https://kusto.test.com",
"https://kusto.test.com/test;fed=true":
"https://kusto.test.com/test",
"https://kusto.test.com:4242;fed=true":
"https://kusto.test.com:4242",
"https://kusto.test.com:4242/;fed=true":
"https://kusto.test.com:4242",
"https://kusto.test.com:4242/test;fed=true":
"https://kusto.test.com:4242/test",
"https://ade.loganalytics.io/subscriptions/da45f7ac-97c0-4fff-ac66-3b6810eb4f65/resourcegroups"
"/some_resource_group/providers/microsoft.operationalinsights/workspaces/some_workspace":
"https://ade.loganalytics.io/subscriptions/da45f7ac-97c0-4fff-ac66-3b6810eb4f65/resourcegroups"
"/some_resource_group/providers/microsoft.operationalinsights/workspaces/some_workspace",
"https://ade.loganalytics.io/subscriptions/da45f7ac-97c0-4fff-ac66-3b6810eb4f65/resourcegroups"
"/some_resource_group/providers/microsoft.operationalinsights/workspaces/some_workspace/":
"https://ade.loganalytics.io/subscriptions/da45f7ac-97c0-4fff-ac66-3b6810eb4f65/resourcegroups"
"/some_resource_group/providers/microsoft.operationalinsights/workspaces/some_workspace",
"https://ade.loganalytics.io:4242/subscriptions/da45f7ac-97c0-4fff-ac66-3b6810eb4f65/resourcegroups"
"/some_resource_group/providers/microsoft.operationalinsights/workspaces/some_workspace/":
"https://ade.loganalytics.io:4242/subscriptions/da45f7ac-97c0-4fff-ac66-3b6810eb4f65/resourcegroups"
"/some_resource_group/providers/microsoft.operationalinsights/workspaces/some_workspace",
"https://ade.loganalytics.io:4242/subscriptions/da45f7ac-97c0-4fff-ac66-3b6810eb4f65/resourcegroups"
"/some_resource_group/providers/microsoft.operationalinsights/workspaces/some_workspace/;fed=true":
"https://ade.loganalytics.io:4242/subscriptions/da45f7ac-97c0-4fff-ac66-3b6810eb4f65"
"/resourcegroups/some_resource_group/providers/microsoft.operationalinsights/workspaces"
"/some_workspace",
"https://kusto.aria.microsoft.com":
"https://kusto.aria.microsoft.com",
"https://kusto.aria.microsoft.com/":
"https://kusto.aria.microsoft.com",
"https://kusto.aria.microsoft.com/;fed=true":
"https://kusto.aria.microsoft.com",
}
for actual_url, expected_url in tests.items():
kcsb = KustoConnectionStringBuilder(actual_url)
assert kcsb.data_source == expected_url
def test_token_provider_auth():
valid_token_provider = lambda: "caller token"
invalid_token_provider = lambda: 12345678
valid_kcsb = KustoConnectionStringBuilder.with_token_provider(
KUSTO_TEST_URI, valid_token_provider)
invalid_kcsb = KustoConnectionStringBuilder.with_token_provider(
KUSTO_TEST_URI, 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 == CallbackTokenProvider.name()
assert str(e.exception).index(str(type(invalid_token_provider()))) > -1
def initialize_kusto_client():
"""initialize kusto client
"""
global KUSTO_INGESTION_CLIENT
if not KUSTO_INGESTION_CLIENT:
kcsb_ingest = KustoConnectionStringBuilder.with_aad_application_key_authentication( \
INGESTION_SERVER_URI, APP_CLIENT_ID, APP_CLIENT_SECRETS, APP_AAD_TENANT_ID)
KUSTO_INGESTION_CLIENT = KustoIngestClient(kcsb_ingest)
logging.info(f"{LOG_MESSAGE_HEADER} Build KUSTO_INGESTION_CLIENT")
else:
logging.info(f"{LOG_MESSAGE_HEADER} KUSTO_INGESTION_CLIENT exist")
def from_dm_kcsb(dm_kcsb: Union[KustoConnectionStringBuilder, str]) -> "ManagedStreamingIngestClient":
"""
Create a ManagedStreamingIngestClient from a KustoConnectionStringBuilder for the dm.
This Connection String is used for the queued ingest client.
This method will infer the engine connection string (by using the same authentication and removing the ingest- prefix)
For advanced use cases, use the constructor directly.
:param dm_kcsb: KustoConnectionStringBuilder for the dm.
:return: ManagedStreamingIngestClient
"""
kcsb = repr(dm_kcsb) if type(dm_kcsb) == KustoConnectionStringBuilder else dm_kcsb
engine_kcsb = KustoConnectionStringBuilder(kcsb.replace("https://ingest-", "https://"))
return ManagedStreamingIngestClient(engine_kcsb, dm_kcsb)
def __init__(self, db_name: str):
"""Initialize a Kusto report DB connector.
Args:
db_name: The Kusto database to connect to.
"""
self.db_name = db_name
ingest_cluster = os.getenv("TEST_REPORT_INGEST_KUSTO_CLUSTER")
tenant_id = os.getenv("TEST_REPORT_AAD_TENANT_ID")
service_id = os.getenv("TEST_REPORT_AAD_CLIENT_ID")
service_key = os.getenv("TEST_REPORT_AAD_CLIENT_KEY")
if not ingest_cluster or not tenant_id or not service_id or not service_key:
raise RuntimeError("Could not load Kusto Credentials from environment")
kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication(ingest_cluster,
service_id,
service_key,
tenant_id)
self._ingestion_client = KustoIngestClient(kcsb)
"""
Kusto performance depends on the work load of cluster, to improve the high availability of test result data service
by hosting a backup cluster, which is optional.
"""
ingest_cluster = os.getenv("TEST_REPORT_INGEST_KUSTO_CLUSTER_BACKUP")
tenant_id = os.getenv("TEST_REPORT_AAD_TENANT_ID_BACKUP")
service_id = os.getenv("TEST_REPORT_AAD_CLIENT_ID_BACKUP")
service_key = os.getenv("TEST_REPORT_AAD_CLIENT_KEY_BACKUP")
if not ingest_cluster or not tenant_id or not service_id or not service_key:
print("Could not load backup Kusto Credentials from environment")
self._ingestion_client_backup = None
else:
kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication(ingest_cluster,
service_id,
service_key,
tenant_id)
self._ingestion_client_backup = KustoIngestClient(kcsb)
def generate_connection_string(
cls, cluster_url: str,
authentication_mode: AuthenticationModeOptions
) -> KustoConnectionStringBuilder:
"""
Generates Kusto Connection String based on given Authentication Mode.
:param cluster_url: Cluster to connect to.
:param authentication_mode: User Authentication Mode, Options: (UserPrompt|ManagedIdentity|AppKey|AppCertificate)
:return: A connection string to be used when creating a Client
"""
# Learn More: For additional information on how to authorize users and apps in Kusto,
# see: https://docs.microsoft.com/azure/data-explorer/manage-database-permissions
if authentication_mode == AuthenticationModeOptions.UserPrompt.name:
# Prompt user for credentials
return KustoConnectionStringBuilder.with_interactive_login(
cluster_url)
elif authentication_mode == AuthenticationModeOptions.ManagedIdentity.name:
# Authenticate using a System-Assigned managed identity provided to an azure service, or using a User-Assigned managed identity.
# For more information, see https://docs.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/overview
return cls.create_managed_identity_connection_string(
cluster_url)
elif authentication_mode == AuthenticationModeOptions.AppKey.name:
# Learn More: For information about how to procure an AAD Application,
# see: https://docs.microsoft.com/azure/data-explorer/provision-azure-ad-app
# TODO (config - optional): App ID & tenant, and App Key to authenticate with
return KustoConnectionStringBuilder.with_aad_application_key_authentication(
cluster_url, os.environ.get("APP_ID"),
os.environ.get("APP_KEY"), os.environ.get("APP_TENANT"))
elif authentication_mode == AuthenticationModeOptions.AppCertificate.name:
return cls.create_application_certificate_connection_string(
cluster_url)
else:
Utils.error_handler(
f"Authentication mode '{authentication_mode}' is not supported"
)
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
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_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(i) for i in kcsb]
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 __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 test_interactive_login():
if not TEST_INTERACTIVE_AUTH:
print(" *** Skipped interactive login Test ***")
return
kcsb = KustoConnectionStringBuilder.with_interactive_login(KUSTO_TEST_URI)
aad_helper = _AadHelper(kcsb)
# should prompt
header = aad_helper.acquire_authorization_header()
assert header is not None
# should not prompt
header = aad_helper.acquire_authorization_header()
assert header is not None
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:
assert kcsb.data_source == "localhost"
assert kcsb.aad_federated_security
assert kcsb.aad_user_id == user
assert kcsb.password == password
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;AAD User ID={0};Password={1};Authority Id=common".format(
user, password)
assert str(
kcsb
) == "Data Source=localhost;AAD Federated Security=True;AAD User ID={0};Password={1};Authority Id=common".format(
user, self.PASSWORDS_REPLACEMENT)
def main():
args = parse_arguments()
# connect to database
kcsb_ingest = KustoConnectionStringBuilder.with_az_cli_authentication(cluster_ingest)
ingest_client = QueuedIngestClient(kcsb_ingest)
date_time = args.datetime
identifier = get_identifier(date_time, args.commit_hash, args.trt_version, args.branch)
try:
result_file = args.report_folder
folders = os.listdir(result_file)
os.chdir(result_file)
tables = [fail_name, memory_name, latency_name, status_name, latency_over_time_name, specs_name, session_name]
table_results = {}
for table_name in tables:
table_results[table_name] = pd.DataFrame()
for model_group in folders:
os.chdir(model_group)
csv_filenames = os.listdir()
for csv in csv_filenames:
table = parse_csv(csv)
if session_name in csv:
table_results[session_name] = table_results[session_name].append(get_session(table, model_group), ignore_index=True)
if specs_name in csv:
table_results[specs_name] = table_results[specs_name].append(get_specs(table, args.branch, args.commit_hash, date_time), ignore_index=True)
if fail_name in csv:
table_results[fail_name] = table_results[fail_name].append(get_failures(table, model_group), ignore_index=True)
if latency_name in csv:
table_results[memory_name] = table_results[memory_name].append(get_memory(table, model_group), ignore_index=True)
table_results[latency_name] = table_results[latency_name].append(get_latency(table, model_group), ignore_index=True)
table_results[latency_over_time_name] = table_results[latency_over_time_name].append(get_latency_over_time(args.commit_hash, args.report_url, args.branch, table_results[latency_name]), ignore_index=True)
if status_name in csv:
table_results[status_name] = table_results[status_name].append(get_status(table, model_group), ignore_index=True)
os.chdir(result_file)
for table in tables:
print('writing ' + table + ' to database')
db_table_name = 'ep_model_' + table
write_table(ingest_client, table_results[table], db_table_name, date_time, identifier)
except BaseException as e:
print(str(e))
sys.exit(1)
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_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, False)
aad_helper.token_provider._init_resources()
try:
aad_helper.acquire_authorization_header()
except KustoAuthenticationError as error:
assert error.authentication_method == UserPassTokenProvider.name()
assert error.authority == "https://login.microsoftonline.com/authorityName"
assert error.kusto_cluster == cluster
assert error.kwargs["username"] == username
assert error.kwargs[
"client_id"] == CloudSettings.DEFAULT_CLOUD.kusto_client_app_id
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