def _run_simple_query_no_close_conn(mock_snowflake, log_tables,
log_resource_usage):
with snowflake_query_tracker(
log_tables=log_tables,
log_resource_usage=log_resource_usage,
log_tables_with_preview=True,
log_tables_with_schema=True,
) as st:
query = "select * from " + TEST_TABLE_NAME
# with self._snowflake_connect() as conn:
conn = _snowflake_connect()
with conn.cursor() as cursor:
cursor.execute(query)
cursor.fetchall()
# we want COMMIT here to have same behavior/amount of queries
# with case when connection is auto-closed (with context manager)
cursor.execute("COMMIT")
assert sorted(st.get_all_tables()) == [TEST_TABLE_NAME]
session_id, query_id = st.get_last_session_with_query_id(many=False)
session_id, query_ids = st.get_last_session_with_query_id(many=True)
session_queries = st.get_all_session_queries().copy()
# 1 session
assert len(session_queries) == 1
for st_session_id, st_query_ids in session_queries.items():
assert len(st_query_ids) == 1
assert (st_session_id, st_query_ids[0]) == (session_id, query_id)
assert (st_session_id, st_query_ids) == (session_id, query_ids)
# query + COMMIT
assert len(mock_snowflake.mock_calls) == 2
assert mock_snowflake.mock_calls[0] == call(query)
if log_resource_usage:
# should be cleaned
assert len(st.get_all_session_queries()) == 0
else:
# on exit from context manager, st.get_all_sessions() shouldn't be affected
# - resources/tables queries shouldn't be tracked anyway
assert len(st.get_all_session_queries()) == len(session_queries)
assert all([
a == b
for a, b in zip(st.get_all_session_queries(), session_queries)
])
return session_id, query_ids
def test_not_tracked_queries(self, mock_snowflake):
with snowflake_query_tracker() as st:
with _snowflake_connect() as conn, conn.cursor() as cursor:
for sql in [
"create table something ()",
"alter table alter column ()",
"alter session set ...",
]:
cursor.execute(sql)
assert len(st.get_all_tables()) == 0
assert len(st.get_all_session_queries()) == 0
assert st.get_last_session_with_query_id(many=True) == (None, [])
assert st.get_last_session_with_query_id(many=False) == (None,
None)
# 1 for automatic "alert session autocommit=false" + 3 queries above
assert len(mock_snowflake.mock_calls) == 4
# should stay the same - no extra queries
assert len(mock_snowflake.mock_calls) == 4