def test_table_fetch_slower_than_timeout(self):
mock_client = MockKustoClient(block=True)
try:
PyKustoClient(mock_client)['test_db']['mock_table']
finally:
# # Return the fetch
mock_client.release()
def test_block_until_fetch_is_done(self):
mock_client = MockKustoClient(block=True, record_metadata=True)
client = PyKustoClient(mock_client)
self.query_in_background(client.get_databases_names)
mock_client.release()
client.wait_for_items()
# Make sure the fetch query was indeed called
assert not mock_client.blocked()
self.assertEqual(self.get_background_query_result(), ('test_db', ))
def test_dir_before_fetch_is_done(self):
mock_client = MockKustoClient(block=True, record_metadata=True)
client = PyKustoClient(mock_client)
self.query_in_background(lambda: dir(client))
# Return the fetch
mock_client.release()
client.wait_for_items()
# Make sure the fetch query was indeed called
assert not mock_client.blocked()
self.assertIn('test_db', self.get_background_query_result())
def test_cross_cluster_join_with_union(self):
client1 = MockKustoClient("https://one.kusto.windows.net")
client2 = MockKustoClient("https://two.kusto.windows.net")
table1 = PyKustoClient(client1)['test_db_1']['test_table_1']
table2 = PyKustoClient(client2)['test_db_2'].get_table('test_table_2_*')
Query(table1).take(5).join(Query(table2).take(6)).on(col.foo).execute()
self.assertEqual(
[RecordedQuery('test_db_1', 'test_table_1 | take 5 | join (union cluster("two.kusto.windows.net").database("test_db_2").table("test_table_2_*") | take 6) on foo')],
client1.recorded_queries,
)
def test_union_column_name_conflict(self):
mock_client = MockKustoClient(
tables_response=mock_tables_response([
('test_table_1', [('foo', _KustoType.STRING), ('bar', _KustoType.INT)]),
('test_table_2', [('foo', _KustoType.BOOL)])
]),
getschema_response=mock_getschema_response([
('foo_string', _KustoType.STRING), ('bar', _KustoType.INT), ('foo_bool', _KustoType.BOOL)
]),
record_metadata=True,
)
db = PyKustoClient(mock_client, fetch_by_default=False)['test_db']
db.blocking_refresh()
table = db.get_table('test_table_*')
table.blocking_refresh() # To trigger name conflict resolution
self.assertEqual(
[
# First trying the usual fetch
RecordedQuery('test_db', '.show database schema | project TableName, ColumnName, ColumnType | limit 10000'),
# Fallback for name conflict resolution
RecordedQuery('test_db', 'union test_table_* | getschema | project ColumnName, DataType | limit 10000')
],
mock_client.recorded_queries,
)
self.assertType(table.foo_string, _StringColumn)
self.assertType(table.bar, _NumberColumn)
self.assertType(table.foo_bool, _BooleanColumn)
def test_column_fetch(self):
mock_kusto_client = MockKustoClient(
columns_response=mock_columns_response([('foo', _KustoType.STRING),
('bar', _KustoType.INT)]),
record_metadata=True,
)
table = PyKustoClient(mock_kusto_client,
fetch_by_default=False)['test_db']['mock_table']
table.blocking_refresh()
# Fetch query
self.assertEqual(
[
RecordedQuery(
'test_db',
'.show table mock_table | project AttributeName, AttributeType | limit 10000'
)
],
mock_kusto_client.recorded_queries,
)
# Dot notation
self.assertEqual(type(table.foo), _StringColumn)
self.assertEqual(type(table.bar), _NumberColumn)
# Bracket notation
self.assertEqual(type(table['foo']), _StringColumn)
self.assertEqual(type(table['bar']), _NumberColumn)
self.assertEqual(type(table['baz']), _AnyTypeColumn)
def test_column_fetch_slow(self):
mock_response_future = Future()
mock_response_future.executed = False
# noinspection PyUnusedLocal
def upon_execute(
query
): # Parameter required since function is passed as Callable[[RecordedQuery], None]
mock_response_future.result()
mock_response_future.executed = True
try:
mock_kusto_client = MockKustoClient(upon_execute=upon_execute,
record_metadata=True)
table = PyKustoClient(
mock_kusto_client,
fetch_by_default=False)['test_db']['mock_table']
table.refresh()
self.assertEqual(type(table['foo']), _AnyTypeColumn)
self.assertEqual(type(table['bar']), _AnyTypeColumn)
self.assertEqual(type(table['baz']), _AnyTypeColumn)
# Make sure above lines were called while the fetch query was still waiting
assert not mock_response_future.executed
finally:
# Return the fetch
mock_response_future.set_result(None)
table.wait_for_items()
# Make sure the fetch query was indeed called
assert mock_response_future.executed
def test_union_wildcard_one_table(self):
mock_kusto_client = MockKustoClient(
tables_response=mock_tables_response([
('test_table_1', [('foo', _KustoType.STRING),
('bar', _KustoType.INT)]),
('other_table_2', [('baz', _KustoType.BOOL)])
]),
record_metadata=True,
)
db = PyKustoClient(mock_kusto_client,
fetch_by_default=False)['test_db']
db.blocking_refresh()
self.assertEqual(
[
RecordedQuery(
'test_db',
'.show database schema | project TableName, ColumnName, ColumnType | limit 10000'
)
],
mock_kusto_client.recorded_queries,
)
table = db.get_table('test_table_*')
self.assertEqual(type(table.foo), _StringColumn)
self.assertEqual(type(table.bar), _NumberColumn)
self.assertEqual(type(table['baz']), _AnyTypeColumn)
def test_execute_already_bound(self):
mock_kusto_client = MockKustoClient()
table = PyKustoClient(mock_kusto_client)['test_db']['mock_table']
self.assertRaises(
RuntimeError("This table is already bound to a query"),
Query(table).take(5).execute, table)
def test_table_fetch(self):
mock_kusto_client = MockKustoClient(
tables_response=mock_tables_response([
('mock_table', [('foo', _KustoType.STRING),
('bar', _KustoType.INT)])
]),
record_metadata=True,
)
db = PyKustoClient(mock_kusto_client,
fetch_by_default=False)['test_db']
db.blocking_refresh()
self.assertEqual(
[
RecordedQuery(
'test_db',
'.show database schema | project TableName, ColumnName, ColumnType | limit 10000'
)
],
mock_kusto_client.recorded_queries,
)
table = db.mock_table
# Table columns
self.assertEqual(type(table.foo), _StringColumn)
self.assertEqual(type(table.bar), _NumberColumn)
self.assertEqual(type(table['baz']), _AnyTypeColumn)
# Bracket notation
self.assertEqual(type(db['other_table']['foo']), _AnyTypeColumn)
# Dot notation error
self.assertRaises(
AttributeError(
"PyKustoClient(test_cluster.kusto.windows.net).Database(test_db) has no attribute 'test_table_1'"
), lambda: db.test_table_1)
def test_union_table(self):
mock_kusto_client = MockKustoClient()
table = PyKustoClient(mock_kusto_client)['test_db'].get_table('test_table1', 'test_table2')
Query(table).take(5).execute()
self.assertEqual(
[RecordedQuery('test_db', 'union test_table1, test_table2 | take 5')],
mock_kusto_client.recorded_queries,
)
def test_single_table_on_execute(self):
mock_kusto_client = MockKustoClient()
table = PyKustoClient(mock_kusto_client)['test_db']['mock_table']
Query().take(5).execute(table)
self.assertEqual(
[RecordedQuery('test_db', 'mock_table | take 5')],
mock_kusto_client.recorded_queries,
)
def test_union_table_with_wildcard(self):
mock_kusto_client = MockKustoClient()
table = PyKustoClient(mock_kusto_client)['test_db']['test_table_*']
Query(table).take(5).execute()
self.assertEqual(
[RecordedQuery('test_db', 'union test_table_* | take 5')],
mock_kusto_client.recorded_queries,
)
def test_union_wildcard_one_table(self):
mock_client = MockKustoClient(record_metadata=True)
db = PyKustoClient(mock_client, fetch_by_default=False)['test_db']
db.blocking_refresh()
self.assertEqual(
[RecordedQuery('test_db', '.show database schema | project TableName, ColumnName, ColumnType | limit 10000')],
mock_client.recorded_queries,
)
table = db.get_table('mock_table_*')
self.assertType(table.foo, _AnyTypeColumn)
self.assertType(table.bar, _AnyTypeColumn)
self.assertType(table['baz'], _BooleanColumn)
def unreliable_mock_kusto_client(
number_of_failures: int,
exception_type: Type[Exception] = KustoServiceError):
TestClient.attempt = 1
def main_response():
if TestClient.attempt <= number_of_failures:
TestClient.attempt += 1
raise exception_type("Mock exception for test", None)
return mock_response(tuple())()
return MockKustoClient(main_response=main_response)
def test_query_before_fetch_returned(self):
mock_response_future = Future()
mock_response_future.returned_queries = []
mock_response_future.called = False
mock_response_future.executed = False
future_called_lock = Lock()
def upon_execute(query):
with future_called_lock:
if mock_response_future.called:
first_run = False
else:
mock_response_future.called = True
first_run = True
if first_run:
mock_response_future.result()
mock_response_future.executed = True
mock_response_future.returned_queries.append(query)
try:
mock_kusto_client = MockKustoClient(upon_execute=upon_execute,
record_metadata=True)
table = PyKustoClient(
mock_kusto_client,
fetch_by_default=False)['test_db']['mock_table']
table.refresh()
# Executing a query in a separate thread, because it is supposed to block until the fetch returns
query_thread = Thread(target=Query(table).take(5).execute)
query_thread.start()
# Make sure above lines were called while the fetch query was still waiting
assert not mock_response_future.executed
finally:
# Return the fetch
mock_response_future.set_result(None)
table.wait_for_items()
query_thread.join()
# Make sure the fetch query was indeed called
assert mock_response_future.executed
# Before the fix the order of returned query was reversed
self.assertEqual(
[
RecordedQuery(
'test_db',
'.show table mock_table | project AttributeName, AttributeType | limit 10000'
),
RecordedQuery('test_db', 'mock_table | take 5'),
],
mock_response_future.returned_queries,
)
def test_request_properties(self):
properties = ClientRequestProperties()
properties.set_option(
ClientRequestProperties.
results_defer_partial_query_failures_option_name, False)
properties.set_parameter('xIntValue', 11)
mock_kusto_client = MockKustoClient()
table = PyKustoClient(mock_kusto_client)['test_db']['mock_table']
Query(table).take(5).execute(properties=properties)
self.assertEqual(
[RecordedQuery('test_db', 'mock_table | take 5', properties)],
mock_kusto_client.recorded_queries)
def test_join_with_table(self):
table = PyKustoClient(
MockKustoClient(columns_response=mock_columns_response([(
'tableStringField',
_KustoType.STRING), (
'numField', _KustoType.INT)])))['test_db']['mock_table']
self.assertEqual(
'mock_table | where numField > 4 | take 5 | join kind=inner (cluster("test_cluster.kusto.windows.net").database("test_db").table("mock_table")) '
'on numField, $left.stringField==$right.tableStringField',
(Query(t).where(t.numField > 4).take(5).join(
Query(table), kind=JoinKind.INNER).on(
t.numField,
(t.stringField, table.tableStringField)).render()))
def test_default_authentication(self):
mock_kusto_client = MockKustoClient()
with patch('pykusto._src.client.PyKustoClient._get_client_for_cluster', lambda s, cluster: mock_kusto_client):
table = PyKustoClient('https://help.kusto.windows.net/')['test_db']['mock_table']
Query().take(5).execute(table)
self.assertIs(
mock_kusto_client,
table._Table__database._Database__client._PyKustoClient__client,
)
self.assertEqual(
[RecordedQuery('test_db', 'mock_table | take 5')],
mock_kusto_client.recorded_queries,
)
def test_autocomplete_with_dot(self):
mock_client = MockKustoClient(
databases_response=mock_databases_response([('test_db', [('mock_table', [('foo', _KustoType.STRING), ('bar.baz', _KustoType.INT)])])]),
)
client = PyKustoClient(mock_client)
client.wait_for_items()
# Table columns
table = client.test_db.mock_table
self.assertType(table.foo, _StringColumn)
self.assertType(table.bar, _AnyTypeColumn)
self.assertType(table['bar.baz'], _NumberColumn)
autocomplete_list = set(dir(client.test_db.mock_table))
self.assertIn('foo', autocomplete_list)
self.assertNotIn('bar.baz', autocomplete_list)
def test_column_fetch_slow(self):
mock_client = MockKustoClient(block=True)
table = PyKustoClient(mock_client, fetch_by_default=False)['test_db']['mock_table']
table.refresh()
mock_client.wait_until_blocked()
self.assertType(table['foo'], _AnyTypeColumn)
self.assertType(table['bar'], _AnyTypeColumn)
self.assertType(table['baz'], _AnyTypeColumn)
# Return the fetch
mock_client.release()
table.wait_for_items()
# Make sure the fetch query was indeed called
assert not mock_client.blocked()
def test_empty_database(self):
mock_client = MockKustoClient(
databases_response=mock_databases_response([
('test_db', [('mock_table', [('foo', _KustoType.STRING), ('bar', _KustoType.INT)])]),
('', [('test_table1', [('foo1', _KustoType.STRING), ('bar1', _KustoType.INT)])])
]),
record_metadata=True,
)
client = PyKustoClient(mock_client)
client.wait_for_items()
self.assertEqual(
[RecordedQuery('', '.show databases schema | project DatabaseName, TableName, ColumnName, ColumnType | limit 100000')],
mock_client.recorded_queries,
)
self.assertType(client.test_db.mock_table.foo, _StringColumn)
def test_query_before_fetch_returned(self):
mock_client = MockKustoClient(block=True, record_metadata=True)
table = PyKustoClient(mock_client, fetch_by_default=False)['test_db']['mock_table']
table.refresh()
mock_client.wait_until_blocked()
mock_client.do_not_block_next_requests()
self.query_in_background(Query(table).take(5).execute)
# Return the fetch
mock_client.release()
table.wait_for_items()
self.get_background_query_result()
# Make sure the fetch query was indeed called
assert not mock_client.blocked()
# Before the fix the order of returned query was reversed
self.assertEqual(
[
RecordedQuery('test_db', '.show table mock_table | project AttributeName, AttributeType | limit 10000'),
RecordedQuery('test_db', 'mock_table | take 5'),
],
mock_client.recorded_queries,
)
def test_to_dataframe(self):
rows = (['foo', 10], ['bar', 20], ['baz', 30])
columns = ('stringField', 'numField')
client = PyKustoClient(
MockKustoClient(
databases_response=mock_databases_response([('test_db', [
('mock_table', [('stringField', _KustoType.STRING),
('numField', _KustoType.INT)])
])]),
main_response=mock_response(rows, columns),
))
client.wait_for_items()
table = client.test_db.mock_table
actual_df = Query(table).take(10).to_dataframe()
expected_df = pd.DataFrame(rows, columns=columns)
self.assertTrue(expected_df.equals(actual_df))
def test_join_chained_on(self):
mock_client = PyKustoClient(
MockKustoClient(columns_response=mock_columns_response([(
'tableStringField',
_KustoType.STRING), ('numField', _KustoType.INT)])))
mock_table = mock_client['test_db']['mock_table']
expected_query = (
'mock_table | where numField > 4 | take 5 | join kind=inner '
'(cluster("test_cluster.kusto.windows.net").database("test_db").table("mock_table")'
' | where numField == 2 | take 6) on numField, $left.stringField==$right.tableStringField'
)
actual_query = (Query(t).where(t.numField > 4).take(5).join(
Query(mock_table).where(mock_table.numField == 2).take(6),
kind=JoinKind.INNER).on(t.numField).on(
(t.stringField, mock_table.tableStringField)).render())
self.assertEqual(expected_query, actual_query)
def test_add_queries_with_table(self):
table = PyKustoClient(
MockKustoClient(columns_response=mock_columns_response([(
'numField', _KustoType.INT)])))['test_db']['mock_table']
query_a = Query(table).where(table.numField > 4)
query_b = Query(t).take(5).take(2).sort_by(t.stringField, Order.ASC,
Nulls.LAST)
query = query_a + query_b
self.assertEqual(
"mock_table | where numField > 4 | take 5 | take 2 | sort by stringField asc nulls last",
query.render(),
)
# make sure the originals didn't change
self.assertEqual(
"mock_table | where numField > 4",
query_a.render(),
)
self.assertEqual(
"mock_table | take 5 | take 2 | sort by stringField asc nulls last",
query_b.render(),
)
def test_database_fetch(self):
mock_kusto_client = MockKustoClient(
databases_response=mock_databases_response([
('test_db', [('mock_table', [('foo', _KustoType.STRING),
('bar', _KustoType.INT)])])
]),
record_metadata=True,
)
client = PyKustoClient(mock_kusto_client)
client.wait_for_items()
self.assertEqual(
[
RecordedQuery(
'',
'.show databases schema | project DatabaseName, TableName, ColumnName, ColumnType | limit 100000'
)
],
mock_kusto_client.recorded_queries,
)
# Table columns
table = client.test_db.mock_table
self.assertEqual(type(table.foo), _StringColumn)
self.assertEqual(type(table.bar), _NumberColumn)
self.assertEqual(type(table['baz']), _AnyTypeColumn)
self.assertEqual(type(client.test_db['other_table']['foo']),
_AnyTypeColumn)
# Various utility methods
db = client.get_database('test_db')
self.assertEqual(type(db), _Database)
self.assertEqual('test_db', db.get_name())
self.assertEqual(('test_db', ), tuple(client.get_databases_names()))
self.assertEqual(('mock_table', 'other_table'),
tuple(client.test_db.get_table_names()))
self.assertEqual(('foo', 'bar', 'baz'),
tuple(client.test_db.mock_table.get_columns_names()))
self.assertTrue({'foo', 'bar'} < set(dir(client.test_db.mock_table)))
self.assertEqual(
'PyKustoClient(test_cluster.kusto.windows.net).Database(test_db).Table(mock_table)',
repr(client.test_db.mock_table))
def test_two_tables_fetch(self):
mock_client = MockKustoClient(record_metadata=True)
db = PyKustoClient(mock_client, fetch_by_default=False)['test_db']
db.blocking_refresh()
self.assertEqual(
[RecordedQuery('test_db', '.show database schema | project TableName, ColumnName, ColumnType | limit 10000')],
mock_client.recorded_queries,
)
# Table columns
self.assertType(db.mock_table.foo, _StringColumn)
self.assertType(db.mock_table.bar, _NumberColumn)
self.assertType(db.mock_table_2['baz'], _BooleanColumn)
self.assertType(db['other_table']['foo'], _AnyTypeColumn)
# Union
table = db.get_table('mock_table', 'mock_table_2')
self.assertType(table.foo, _StringColumn)
self.assertType(table.bar, _NumberColumn)
self.assertType(table.baz, _BooleanColumn)
# Wildcard
table = db.get_table('mock_*')
self.assertType(table.foo, _StringColumn)
self.assertType(table.bar, _NumberColumn)
self.assertType(table.baz, _BooleanColumn)
def test_client_not_fetched(self):
client = PyKustoClient(MockKustoClient(), fetch_by_default=False)
self.assertEqual(frozenset(), set(client.get_databases_names()))
self.assertEqual(frozenset(), set(client.get_databases()))
def test_exception_while_fetching(self):
client = PyKustoClient(MockKustoClient(databases_response=self.raise_mock_exception))
self.assertRaises(
Exception("Mock exception"),
lambda: set(client.get_databases_names()),
)