def _log_snowflake_resources(
query_text,
snowflake_conn_id,
session_id=None,
warehouse=None,
database=None,
role=None,
schema=None,
account=None,
):
hook = SnowflakeHook(
snowflake_conn_id,
warehouse=warehouse,
account=account,
database=database,
role=role,
schema=schema,
)
conn = hook.get_uri()
conn_params = hook._get_conn_params()
log_snowflake_resource_usage(
query_text,
database=hook.database,
user=conn_params["user"],
connection_string=conn,
session_id=session_id,
)
def update_customers_with_monitoring(**kwargs):
snowflake_hook = DbndSnowflakeHook(snowflake_conn_id=SNOWFLAKE_CONNECTION_ID)
session_id, query_ids = snowflake_hook.run(update_query)
log_snowflake_resource_usage(
database=database,
connection_string=snowflake_hook.get_uri(),
session_id=session_id,
query_ids=query_ids,
key="DbndSnowflakeHook_get_records",
)
def process_customers_with_monitoring(**kwargs):
snowflake_hook = DbndSnowflakeHook(snowflake_conn_id=SNOWFLAKE_CONNECTION_ID)
customers, session_id, query_id = snowflake_hook.get_records(select_query)
log_snowflake_resource_usage(
database=database,
connection_string=snowflake_hook.get_uri(),
session_id=session_id,
query_ids=[query_id],
key="DbndSnowflakeHook_get_records",
)
# Process records - Same code
process_records(customers)
def execute(self, context):
hook = self.get_hook()
log_snowflake_resource_usage(
database=hook.database,
connection_string=hook.get_uri(),
query_ids=[self.query_id],
session_id=int(self.session_id) if self.session_id else None,
key=self.key,
history_window=self.history_window,
query_history_result_limit=self.query_history_result_limit,
retries=self.retries,
retry_pause=self.retry_pause,
raise_on_error=self.raise_on_error,
)
def log_all_snowflake_resource_usage(cd, reset=True, **kwargs):
# type: (ConnectionData, bool, ...) -> Optional[Dict[int, List[str]]]
# copy to prevent resource usage queries to appear in current history
history = cd.get_all_session_queries()
for session_id, query_ids in history.items():
log_snowflake_resource_usage(
database=cd.database,
connection_string=cd.connection,
session_id=session_id,
query_ids=query_ids,
**kwargs
)
if reset:
cd.reset_session_queries()
return history
def execute(self, context):
hook = SnowflakeHook(snowflake_conn_id=snowflake_conn_id)
session_id, query_ids = snowflake_run(hook, select_query)
connection_string = hook.get_uri()
log_snowflake_table(
table_name=table,
connection_string=connection_string,
database=database,
schema=schema,
key=f"example1.{table}",
with_preview=False,
raise_on_error=False,
)
log_snowflake_resource_usage(
database=database,
key=f"example1.{session_id}{query_ids[0]}",
connection_string=connection_string,
query_ids=query_ids,
session_id=int(session_id),
raise_on_error=True,
)