def run_query(self,
query: str,
database: str,
options: Optional[Dict] = None) -> KustoResponseDataSetV2:
"""
Run KQL query using provided configuration, and return
`azure.kusto.data.response.KustoResponseDataSet` instance.
If query is unsuccessful AirflowException is raised.
:param query: KQL query to run
:type query: str
:param database: Database to run the query on.
:type database: str
:param options: Optional query options. See:
https://docs.microsoft.com/en-us/azure/kusto/api/netfx/request-properties#list-of-clientrequestproperties
:type options: dict
:return: dict
"""
properties = ClientRequestProperties()
if options:
for k, v in options.items():
properties.set_option(k, v)
try:
return self.connection.execute(database,
query,
properties=properties)
except KustoServiceError as error:
raise AirflowException(f'Error running Kusto query: {error}')
def test_properties(self):
"""positive tests"""
defer = False
timeout = timedelta(seconds=10)
crp = ClientRequestProperties()
crp.set_option(
ClientRequestProperties.
results_defer_partial_query_failures_option_name, defer)
crp.set_option(ClientRequestProperties.request_timeout_option_name,
timeout)
result = crp.to_json()
assert '"{0}": false'.format(
crp.results_defer_partial_query_failures_option_name) in result
assert '"{0}": "0:00:10"'.format(
ClientRequestProperties.request_timeout_option_name) in result
assert crp.client_request_id is None
assert crp.application is None
assert crp.user is None
crp.client_request_id = "CRID"
assert crp.client_request_id == "CRID"
crp.application = "myApp"
assert crp.application == "myApp"
crp.user = "******"
assert crp.user == "myUser"
def test_properties(self):
"""positive tests"""
defer = False
timeout = timedelta(seconds=10)
crp = ClientRequestProperties()
crp.set_option(ClientRequestProperties.OptionDeferPartialQueryFailures,
defer)
crp.set_option(ClientRequestProperties.OptionServerTimeout, timeout)
result = crp.to_json()
assert '"{0}": false'.format(
crp.OptionDeferPartialQueryFailures) in result
assert '"{0}": "0:00:10"'.format(
ClientRequestProperties.OptionServerTimeout) in result
def test_run_query(self, mock_execute):
mock_execute.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'})))
hook = AzureDataExplorerHook(
azure_data_explorer_conn_id=ADX_TEST_CONN_ID)
hook.run_query('Database',
'Logs | schema',
options={'option1': 'option_value'})
properties = ClientRequestProperties()
properties.set_option('option1', 'option_value')
assert mock_execute.called_with('Database',
'Logs | schema',
properties=properties)
def test_partial_results(self, mock_post):
"""Tests partial results."""
client = KustoClient("https://somecluster.kusto.windows.net")
query = """set truncationmaxrecords = 5;
range x from 1 to 10 step 1"""
properties = ClientRequestProperties()
properties.set_option(
ClientRequestProperties.OptionDeferPartialQueryFailures, False)
self.assertRaises(KustoServiceError, client.execute_query,
"PythonTest", query, properties)
properties.set_option(
ClientRequestProperties.OptionDeferPartialQueryFailures, True)
response = client.execute_query("PythonTest", query, properties)
self.assertEqual(response.errors_count, 1)
self.assertIn("E_QUERY_RESULT_SET_TOO_LARGE",
response.get_exceptions()[0])
self.assertEqual(len(response), 3)
results = list(response.primary_results[0])
self.assertEqual(len(results), 5)
self.assertEqual(results[0]["x"], 1)
def test_partial_results(self, mock_post):
"""Tests partial results."""
client = KustoClient("https://somecluster.kusto.windows.net")
query = """set truncationmaxrecords = 5;
range x from 1 to 10 step 1"""
properties = ClientRequestProperties()
properties.set_option(
ClientRequestProperties.
results_defer_partial_query_failures_option_name, False)
self.assertRaises(KustoServiceError, client.execute_query,
"PythonTest", query, properties)
properties.set_option(
ClientRequestProperties.
results_defer_partial_query_failures_option_name, True)
response = client.execute_query("PythonTest", query, properties)
assert response.errors_count == 1
assert "E_QUERY_RESULT_SET_TOO_LARGE" in response.get_exceptions()[0]
assert len(response) == 3
results = list(response.primary_results[0])
assert len(results) == 5
assert results[0]["x"] == 1
##################
### EXCEPTIONS ###
##################
# Query is too big to be executed
query = "StormEvents"
try:
response = client.execute(db, query)
except KustoServiceError as error:
print("2. Error:", error)
print("2. Is semantic error:", error.is_semantic_error())
print("2. Has partial results:", error.has_partial_results())
print("2. Result size:", len(error.get_partial_results()))
properties = ClientRequestProperties()
properties.set_option(properties.OptionDeferPartialQueryFailures, True)
properties.set_option(properties.OptionServerTimeout,
timedelta(seconds=8 * 60))
response = client.execute(db, query, properties=properties)
print("3. Response error count: ", response.errors_count)
print("3. Exceptions:", response.get_exceptions())
print("3. Result size:", len(response.primary_results))
# Query has semantic error
query = "StormEvents | where foo = bar"
try:
response = client.execute(db, query)
except KustoServiceError as error:
print("4. Error:", error)
print("4. Is semantic error:", error.is_semantic_error())
print("4. Has partial results:", error.has_partial_results())
##################
### EXCEPTIONS ###
##################
# Query is too big to be executed
query = "StormEvents"
try:
response = client.execute(db, query)
except KustoServiceError as error:
print("2. Error:", error)
print("2. Is semantic error:", error.is_semantic_error())
print("2. Has partial results:", error.has_partial_results())
print("2. Result size:", len(error.get_partial_results()))
properties = ClientRequestProperties()
properties.set_option(
properties.results_defer_partial_query_failures_option_name, True)
properties.set_option(properties.request_timeout_option_name,
timedelta(seconds=8 * 60))
response = client.execute(db, query, properties=properties)
print("3. Response error count: ", response.errors_count)
print("3. Exceptions:", response.get_exceptions())
print("3. Result size:", len(response.primary_results))
# Query has semantic error
query = "StormEvents | where foo = bar"
try:
response = client.execute(db, query)
except KustoServiceError as error:
print("4. Error:", error)
print("4. Is semantic error:", error.is_semantic_error())
print("4. Has partial results:", error.has_partial_results())
##################
### EXCEPTIONS ###
##################
# Query is too big to be executed
query = "StormEvents"
try:
response = client.execute(db, query)
except KustoServiceError as error:
print("2. Error:", error)
print("2. Is semantic error:", error.is_semantic_error())
print("2. Has partial results:", error.has_partial_results())
print("2. Result size:", len(error.get_partial_results()))
properties = ClientRequestProperties()
properties.set_option(properties.OptionDeferPartialQueryFailures, True)
response = client.execute(db, query, properties=properties)
print("3. Response error count: ", response.errors_count)
print("3. Exceptions:", response.get_exceptions())
print("3. Result size:", len(response.primary_results))
# Query has semantic error
query = "StormEvents | where foo = bar"
try:
response = client.execute(db, query)
except KustoServiceError as error:
print("4. Error:", error)
print("4. Is semantic error:", error.is_semantic_error())
print("4. Has partial results:", error.has_partial_results())
client = KustoClient("https://kustolab.kusto.windows.net")
kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication(
cluster, client_id, client_secret, tenant_id)
kusto_client = KustoClient(kcsb)
######################################################
## QUERY ##
######################################################
db_name = os.environ.get("DATABASE_NAME")
test_id = os.environ.get("TEST_ID", str(uuid.uuid4()))
query_consistency = os.environ.get("QUERY_CONSISTENCY", "weakconsistency")
request_properties = ClientRequestProperties()
request_properties.set_option("queryconsistency", query_consistency)
instrumentation_key = os.environ.get("APPINSIGHTS_INSTRUMENTATIONKEY")
telemetry_client = None
if instrumentation_key:
telemetry_client = TelemetryClient(instrumentation_key)
print("Test run for '{}' started.".format(test_id))
def execute_query(raw_query):
query = "{} //TEST_ID={}".format(raw_query, test_id)
response = kusto_client.execute(db_name, query, request_properties)
if response.errors_count > 0:
if telemetry_client: