def test_null_values_in_data(self, mock_post):
"""Tests response with null values in non nullable column types"""
client = KustoClient(self.HOST)
query = "PrimaryResultName"
response = client.execute_query("PythonTest", query)
assert response is not None
def test_identifiable_os(self, mock_post):
"""Tests identifiable OS doesn't fail when composing its socket options"""
with patch.object(sys, "platform", "win32"):
client = KustoClient("https://somecluster.kusto.windows.net")
query = """print dynamic(123)"""
row = client.execute_query("PythonTest", query).primary_results[0].rows[0]
assert isinstance(row[0], int)
def setup_class(cls):
cls.engine_cs = os.environ.get("ENGINE_CONNECTION_STRING") or ""
cls.ai_engine_cs = os.environ.get("APPLICATION_INSIGHTS_ENGINE_CONNECTION_STRING") or ""
cls.app_id = os.environ.get("APP_ID")
cls.app_key = os.environ.get("APP_KEY")
cls.auth_id = os.environ.get("AUTH_ID")
cls.test_db = os.environ.get("TEST_DATABASE")
cls.ai_test_db = os.environ.get("APPLICATION_INSIGHTS_TEST_DATABASE") # name of e2e database could be changed
if not all([cls.engine_cs, cls.test_db, cls.ai_engine_cs, cls.ai_test_db]):
pytest.skip("E2E environment is missing")
# Init clients
cls.streaming_test_table = "BigChunkus"
cls.streaming_test_table_query = cls.streaming_test_table + " | order by timestamp"
cls.ai_test_table_cmd = ".show tables"
cls.client = KustoClient(cls.engine_kcsb_from_env())
cls.ai_client = KustoClient(cls.engine_kcsb_from_env(True))
cls.input_folder_path = cls.get_file_path()
with open(os.path.join(cls.input_folder_path, "big.json")) as f:
cls.test_streaming_data = json.load(f)
def test_null_values_in_data(self, mock_post):
"""Tests response with null values in non nullable column types"""
client = KustoClient("https://somecluster.kusto.windows.net")
query = "PrimaryResultName"
response = client.execute_query("PythonTest", query)
assert response is not None
def test_dynamic(self, mock_post):
"""Tests dynamic responses."""
client = KustoClient(self.HOST)
query = """print dynamic(123), dynamic("123"), dynamic("test bad json"),"""
""" dynamic(null), dynamic('{"rowId":2,"arr":[0,2]}'), dynamic({"rowId":2,"arr":[0,2]})"""
row = client.execute_query("PythonTest",
query).primary_results[0].rows[0]
self._assert_dynamic_response(row)
class KustoStreamingIngestClient(BaseIngestClient):
"""Kusto streaming ingest client for Python.
KustoStreamingIngestClient works with both 2.x and 3.x flavors of Python.
All primitive types are supported.
Tests are run using pytest.
"""
def __init__(self, kcsb: Union[KustoConnectionStringBuilder, str]):
"""Kusto Streaming Ingest Client constructor.
:param KustoConnectionStringBuilder kcsb: The connection string to initialize KustoClient.
"""
self._kusto_client = KustoClient(kcsb)
def set_proxy(self, proxy_url: str):
self._kusto_client.set_proxy(proxy_url)
def ingest_from_file(self, file_descriptor: Union[FileDescriptor, str], ingestion_properties: IngestionProperties) -> IngestionResult:
"""Ingest from local files.
:param file_descriptor: a FileDescriptor to be ingested.
:param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties.
"""
stream_descriptor = StreamDescriptor.from_file_descriptor(file_descriptor)
with stream_descriptor.stream:
return self.ingest_from_stream(stream_descriptor, ingestion_properties)
def ingest_from_stream(self, stream_descriptor: Union[StreamDescriptor, IO[AnyStr]], ingestion_properties: IngestionProperties) -> IngestionResult:
"""Ingest from io streams.
:param azure.kusto.ingest.StreamDescriptor stream_descriptor: An object that contains a description of the stream to
be ingested.
:param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties.
"""
return self._ingest_from_stream_with_client_request_id(stream_descriptor, ingestion_properties, None)
def _ingest_from_stream_with_client_request_id(
self, stream_descriptor: Union[StreamDescriptor, IO[AnyStr]], ingestion_properties: IngestionProperties, client_request_id: Optional[str]
) -> IngestionResult:
stream_descriptor = BaseIngestClient._prepare_stream(stream_descriptor, ingestion_properties)
additional_properties = None
if client_request_id:
additional_properties = ClientRequestProperties()
additional_properties.client_request_id = client_request_id
self._kusto_client.execute_streaming_ingest(
ingestion_properties.database,
ingestion_properties.table,
stream_descriptor.stream,
ingestion_properties.format.name,
additional_properties,
mapping_name=ingestion_properties.ingestion_mapping_reference,
)
return IngestionResult(IngestionStatus.SUCCESS, ingestion_properties.database, ingestion_properties.table, stream_descriptor.source_id)
def test_admin_then_query(self, mock_post):
"""Tests admin then query."""
client = KustoClient("https://somecluster.kusto.windows.net")
query = ".show tables | project DatabaseName, TableName"
response = client.execute_mgmt("PythonTest", query)
assert response.errors_count == 0
assert len(response) == 4
results = list(response.primary_results[0])
assert len(results) == 2
assert response[0].table_kind == WellKnownDataSet.PrimaryResult
assert response[1].table_kind == WellKnownDataSet.QueryProperties
assert response[2].table_kind == WellKnownDataSet.QueryCompletionInformation
assert response[3].table_kind == WellKnownDataSet.TableOfContents
def test_sanity_control_command(self, mock_post):
"""Tests contol command."""
client = KustoClient("https://somecluster.kusto.windows.net")
response = client.execute_mgmt("NetDefaultDB", ".show version")
assert len(response) == 1
primary_table = response.primary_results[0]
row_count = 0
for _ in primary_table:
row_count += 1
assert row_count == 1
result = primary_table[0]
assert result["BuildVersion"] == "1.0.6693.14577"
assert result["BuildTime"] == datetime(year=2018, month=4, day=29, hour=8, minute=5, second=54, tzinfo=UTC)
assert result["ServiceType"] == "Engine"
assert result["ProductVersion"] == "KustoMain_2018.04.29.5"
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
def test_partial_results(self, mock_post):
"""Tests partial results."""
client = KustoClient(self.HOST)
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)
self._assert_partial_results_response(response)
def authenticate_to_kusto(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 = KustoClient(kcsb)
return kusto_client
def test_null_values_in_data(self, mock_post, method):
"""Tests response with null values in non nullable column types"""
client = KustoClient(self.HOST)
query = "PrimaryResultName"
response = method.__call__(client, "PythonTest", query)
assert response is not None
def test_sanity_data_frame(self, mock_post, method):
"""Tests KustoResponse to pandas.DataFrame."""
client = KustoClient(self.HOST)
response = method.__call__(client, "PythonTest", "Deft")
data_frame = dataframe_from_result_table(
get_response_first_primary_result(response))
self._assert_sanity_data_frame_response(data_frame)
def test_partial_results(self, mock_post, method):
"""Tests partial results."""
client = KustoClient(self.HOST)
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)
with pytest.raises(KustoMultiApiError) as e:
response = method.__call__(client,
"PythonTest",
query,
properties=properties)
if type(response) == KustoStreamingResponseDataSet:
results = list(get_response_first_primary_result(response))
errors = e.value.get_api_errors()
assert len(errors) == 1
assert errors[0].code == "LimitsExceeded"
properties.set_option(
ClientRequestProperties.
results_defer_partial_query_failures_option_name, True)
response = method.__call__(client,
"PythonTest",
query,
properties=properties)
self._assert_partial_results_response(response)
def test_partial_results(self, mock_post, mock_aad):
""" Tests partial results. """
client = KustoClient("https://somecluster.kusto.windows.net")
query = """\
set truncationmaxrecords = 1;
range x from 1 to 2 step 1 """
self.assertRaises(KustoServiceError, client.execute_query,
"PythonTest", query)
response = client.execute_query("PythonTest",
query,
accept_partial_results=True)
self.assertTrue(response.has_exceptions())
self.assertEqual(response.get_exceptions()[0]['error']['code'],
'LimitsExceeded')
self.assertEqual(response.get_table_count(), 5)
results = list(response.iter_all())
self.assertEqual(len(results), 1)
self.assertEqual(results[0]['x'], 1)
def test_dynamic(self, mock_post, method):
"""Tests dynamic responses."""
client = KustoClient(self.HOST)
query = """print dynamic(123), dynamic("123"), dynamic("test bad json"),"""
""" dynamic(null), dynamic('{"rowId":2,"arr":[0,2]}'), dynamic({"rowId":2,"arr":[0,2]})"""
row = get_table_first_row(
get_response_first_primary_result(
method.__call__(client, "PythonTest", query)))
self._assert_dynamic_response(row)
def test_identifiable_os(self, mock_post, method):
"""Tests identifiable OS doesn't fail when composing its socket options"""
with patch.object(sys, "platform", "win32"):
client = KustoClient("https://somecluster.kusto.windows.net")
query = """print dynamic(123)"""
row = get_table_first_row(
get_response_first_primary_result(
method.__call__(client, "PythonTest", query)))
assert isinstance(row[0], int)
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 get_all_db_table_records_count(kusto_client: KustoClient,
database: str,
table_prefix: str = ''):
# Count Records from CO Table and TEMP Tables in All Databases
query_filter = f"union database(\"*\").CO*|union database(\"*\").TEMP*|count"
print(f"query on database: {database} by query_filter: {query_filter}")
response = kusto_client.execute(database, query_filter)
result = response.primary_results[0]
records_count = result.to_dict().get('data')[0]
return records_count.get('Count')
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 __init__(self,
kusto_cluster,
client_id=None,
client_secret=None,
username=None,
password=None,
authority=None):
"""
Kusto Client constructor.
Parameters
----------
kusto_cluster : str
Kusto cluster endpoint. Example: https://ingest-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
version : 'v1', optional
REST API version, defaults to v1.
authority : 'microsoft.com', optional
In case your tenant is not microsoft please use this param.
"""
self._kusto_client = KustoClient(kusto_cluster,
client_id=client_id,
client_secret=client_secret,
username=username,
password=password,
authority=authority)
self._last_time_refreshed_containers = datetime.min
self._temp_storage_objects = None
self._last_time_refreshed_queues = datetime.min
self._queues = None
self._last_time_refreshed_token = datetime.min
self._kusto_token = None
def test_dynamic(self, mock_post):
"""Tests dynamic responses."""
client = KustoClient("https://somecluster.kusto.windows.net")
query = """print dynamic(123), dynamic("123"), dynamic("test bad json"),"""
""" dynamic(null), dynamic('{"rowId":2,"arr":[0,2]}'), dynamic({"rowId":2,"arr":[0,2]})"""
row = client.execute_query("PythonTest", query).primary_results[0].rows[0]
assert isinstance(row[0], int)
assert row[0] == 123
assert isinstance(row[1], str)
assert row[1] == "123"
assert isinstance(row[2], str)
assert row[2] == "test bad json"
assert row[3] is None
assert isinstance(row[4], str)
assert row[4] == '{"rowId":2,"arr":[0,2]}'
assert isinstance(row[5], dict)
assert row[5] == {"rowId": 2, "arr": [0, 2]}
def test_custom_request_id(self, mock_post, method):
"""Test query V2."""
client = KustoClient(self.HOST)
properties = ClientRequestProperties()
request_id = "test_request_id"
properties.client_request_id = request_id
response = method.__call__(client,
"PythonTest",
"Deft",
properties=properties)
self._assert_sanity_query_response(response)
self._assert_client_request_id(mock_post.call_args.kwargs,
value=request_id)
def test_sanity_control_command(self, mock_post, mock_aad):
""" Tests contol command. """
client = KustoClient("https://somecluster.kusto.windows.net")
response = client.execute_mgmt("NetDefaultDB", ".show version")
self.assertEqual(response.get_table_count(), 1)
row_count = 0
for _ in response.iter_all():
row_count += 1
self.assertEqual(row_count, 1)
result = list(response.iter_all())[0]
self.assertEqual(result['BuildVersion'], "1.0.6693.14577")
self.assertEqual(
result['BuildTime'],
datetime(year=2018,
month=4,
day=29,
hour=8,
minute=5,
second=54,
tzinfo=tzutc()))
self.assertEqual(result['ServiceType'], "Engine")
self.assertEqual(result['ProductVersion'], "KustoMain_2018.04.29.5")
def test_proxy_token_providers(self, mock_get, proxy_kcsb):
"""Test query V2."""
proxy = "https://my_proxy.sample"
kcsb, auth_supports_proxy = proxy_kcsb
client = KustoClient(kcsb)
client.set_proxy(proxy)
assert client._proxy_url == proxy
expected_dict = {"http": proxy, "https": proxy}
if not auth_supports_proxy:
return
assert client._aad_helper.token_provider._proxy_dict == expected_dict
assert client._session.proxies == expected_dict
CloudSettings._cloud_cache.clear()
client._aad_helper.token_provider._init_resources()
mock_get.assert_called_with(
"https://somecluster.kusto.windows.net/v1/rest/auth/metadata",
proxies=expected_dict)
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 get_dataset():
in_use = microsoft_config
kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication(
in_use['cluster'], in_use['client_id'], in_use["client_secret"], in_use["authority_id"])
query = '''
RawEventsAzCli
| where EventTimestamp > ago(2h) and EventTimestamp < ago(1h)
// Remove pre 2.0.28 telemetry as it has a different schema
| where toint(split(tostring(parse_json(Properties).["context.default.vs.core.telemetryapi.productversion"]), '.')[2]) > 28
| where tostring(parse_json(Properties).["reserved.datamodel.action.result"]) == "Success"
| sort by UserId
| project UserId,command=parse_json(Properties)["context.default.azurecli.rawcommand"],
params = parse_json(Properties).["context.default.azurecli.params"],
os_type = tostring(parse_json(Properties).["context.default.vs.core.os.type"]),
source = tostring(parse_json(Properties).["context.default.azurecli.source"]),
start_time_str= todatetime(parse_json(Properties).["context.default.azurecli.starttime"]),
end_time_str = todatetime(parse_json(Properties).["context.default.azurecli.endtime"])
| take 500
'''
client = KustoClient(kcsb)
response = client.execute("AzureCli", query)
dataset = {}
for row in response.primary_results[0]:
uid, command, params, os_type, source, begin, end = row
params = sort_params(params)
if uid not in dataset:
dataset[uid] = [(command, params, os_type, source, begin, end)]
else:
last_cmd, last_params, *_ = dataset[uid][-1]
last_params = sort_params(last_params)
if last_cmd == command and last_params == params:
continue # duplicate
dataset[uid].append(
(command, params, os_type, source, begin, end))
return dataset
def setup_class(cls):
# DM CS can be composed from engine CS
cls.engine_cs = os.environ.get("ENGINE_CONNECTION_STRING")
cls.dm_cs = os.environ.get(
"DM_CONNECTION_STRING") or cls.engine_cs.replace(
"//", "//ingest-")
cls.app_id = os.environ.get("APP_ID")
cls.app_key = os.environ.get("APP_KEY")
cls.auth_id = os.environ.get("AUTH_ID")
cls.test_db = os.environ.get("TEST_DATABASE")
if not all([
cls.engine_cs, cls.dm_cs, cls.app_id, cls.app_key, cls.auth_id,
cls.test_db
]):
raise unittest.SkipTest("E2E environment is missing")
# Init clients
python_version = "_".join([str(v) for v in sys.version_info[:3]])
cls.test_table = "python_test_{0}_{1}_{2}".format(
python_version, str(int(time.time())), random.randint(1, 100000))
cls.client = KustoClient(cls.engine_kcsb_from_env())
cls.ingest_client = KustoIngestClient(cls.dm_kcsb_from_env())
cls.streaming_ingest_client = KustoStreamingIngestClient(
cls.engine_kcsb_from_env())
cls.input_folder_path = cls.get_file_path()
cls.csv_file_path = os.path.join(cls.input_folder_path, "dataset.csv")
cls.tsv_file_path = os.path.join(cls.input_folder_path, "dataset.tsv")
cls.zipped_csv_file_path = os.path.join(cls.input_folder_path,
"dataset.csv.gz")
cls.json_file_path = os.path.join(cls.input_folder_path,
"dataset.json")
cls.zipped_json_file_path = os.path.join(cls.input_folder_path,
"dataset.jsonz.gz")
cls.current_count = 0
cls.client.execute(
cls.test_db,
".create table {0} (rownumber: int, rowguid: string, xdouble: real, xfloat: real, xbool: bool, xint16: int, xint32: int, xint64: long, xuint8: long, xuint16: long, xuint32: long, xuint64: long, xdate: datetime, xsmalltext: string, xtext: string, xnumberAsText: string, xtime: timespan, xtextWithNulls: string, xdynamicWithNulls: dynamic)"
.format(cls.test_table),
)
cls.client.execute(
cls.test_db,
".create table {0} ingestion json mapping 'JsonMapping' {1}".
format(cls.test_table, cls.test_table_json_mapping_reference()))
def _set_client(self):
self.client = KustoClient(
kusto_cluster=self.cluster_url_template.format(
self._parsed_conn.get('cluster')),
client_id=self._parsed_conn.get('clientid'),
client_secret=self._parsed_conn.get('clientsecret'),
username=self._parsed_conn.get('username'),
password=self._parsed_conn.get('password'),
authority=self._parsed_conn.get('tenant'))
# patch that replace the authetication helper
my_aad_helper = _MyAadHelper(
kusto_cluster=self.cluster_url_template.format(
self._parsed_conn.get('cluster')),
client_id=self._parsed_conn.get('clientid'),
client_secret=self._parsed_conn.get('clientsecret'),
username=self._parsed_conn.get('username'),
password=self._parsed_conn.get('password'),
authority=self._parsed_conn.get('tenant'))
self.client._aad_helper = my_aad_helper
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
