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 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 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 __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 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_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_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_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 __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 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)
)
# In case you want to authenticate with AAD application certificate.
filename = "path to a PEM certificate"
with open(filename, "r") as pem_file:
PEM = pem_file.read()
thumbprint = "certificate's thumbprint"
kcsb = KustoConnectionStringBuilder.with_aad_application_certificate_authentication(
cluster, client_id, PEM, thumbprint, authority_id
)
# In case you want to authenticate with AAD username and password
username = "******"
password = "******"
kcsb = KustoConnectionStringBuilder.with_aad_user_password_authentication(cluster, username, password, authority_id)
# In case you want to authenticate with AAD device code.
# Please note that if you choose this option, you'll need to autenticate for every new instance that is initialized.
# It is highly recommended to create one instance and use it for all of your queries.
kcsb = KustoConnectionStringBuilder.with_aad_device_authentication(cluster)
# The authentication method will be taken from the chosen KustoConnectionStringBuilder.
client = KustoIngestClient(kcsb)
# there are more options for authenticating - see azure-kusto-data samples
##################################################################
## INGESTION ##
##################################################################
