diff --git a/src/snowflake/snowpark/profiler.py b/src/snowflake/snowpark/profiler.py index 5aba9357f70..4e2c36f135d 100644 --- a/src/snowflake/snowpark/profiler.py +++ b/src/snowflake/snowpark/profiler.py @@ -1,6 +1,7 @@ # # Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved. # +import re from contextlib import contextmanager from typing import List, Optional @@ -109,8 +110,11 @@ def disable_profiler(self): self.session.sql(self.disable_profiler_sql).collect() def _get_last_query_id(self): + pattern = r"WITH\s+.*?\s+AS\s+PROCEDURE\s+.*?\s+CALL\s+.*" for query in self.query_history.queries[::-1]: - if query.sql_text.startswith("CALL"): + if query.sql_text.startswith("CALL") or re.match( + pattern, query.sql_text, re.DOTALL + ): return query.query_id return None diff --git a/tests/integ/test_profiler.py b/tests/integ/test_profiler.py index 24eca44b170..bcc5a46d803 100644 --- a/tests/integ/test_profiler.py +++ b/tests/integ/test_profiler.py @@ -80,3 +80,21 @@ def single_value_sp(session: snowflake.snowpark.Session) -> str: session.register_profiler_modules([]) assert res is not None assert "Modules Profiled" in res + + +def test_anonymous_procedure(session, db_parameters): + def single_value_sp(session: snowflake.snowpark.Session) -> str: + return "success" + + single_value_sp = session.sproc.register(single_value_sp, anonymous=True) + session.register_profiler_modules(["table_sp"]) + with profiler( + stage=f"{db_parameters['database']}.{db_parameters['schema']}.{tmp_stage_name}", + active_profiler="LINE", + session=session, + ): + single_value_sp() + res = session.show_profiles() + session.register_profiler_modules([]) + assert res is not None + assert "Modules Profiled" in res