Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-1663726 make session config updates thread safe #2302

Merged
Changes from 1 commit
Commits
Show all changes
72 commits
Select commit Hold shift + click to select a range
56fb566
init
sfc-gh-aalam Sep 11, 2024
66003d1
make udf/sproc related files thread-safe
sfc-gh-aalam Sep 11, 2024
0e58205
Merge branch 'main' into aalam-SNOW-1418523-make-udf-sproc-thread-safe
sfc-gh-aalam Sep 11, 2024
e75dde1
init
sfc-gh-aalam Sep 11, 2024
68a8c1c
make query listener thread-safe
sfc-gh-aalam Sep 11, 2024
31a5734
Fix query_tag and last_action_id
sfc-gh-aalam Sep 11, 2024
b4dadda
core updates done
sfc-gh-aalam Sep 11, 2024
b8c6496
Add tests
sfc-gh-aalam Sep 12, 2024
f39837e
Fix local tests
sfc-gh-aalam Sep 12, 2024
31a196f
Merge branch 'main' into aalam-SNOW-1418523-make-analyzer-server_conn…
sfc-gh-aalam Sep 12, 2024
723bdf7
Merge branch 'aalam-SNOW-1418523-make-internal-session-variables-thre…
sfc-gh-aalam Sep 12, 2024
37c0419
add file IO tests
sfc-gh-aalam Sep 12, 2024
8a2d433
Merge branch 'aalam-SNOW-1418523-concurrent-file-operations' into aal…
sfc-gh-aalam Sep 12, 2024
a083989
make session._runtime_version_from_requirement safe
sfc-gh-aalam Sep 13, 2024
947d384
add sp/udf concurrent tests
sfc-gh-aalam Sep 13, 2024
fd51720
fix broken test
sfc-gh-aalam Sep 13, 2024
3077853
add udtf/udaf tests
sfc-gh-aalam Sep 13, 2024
65c3186
fix broken test
sfc-gh-aalam Sep 13, 2024
94412cf
sql_simplifier, cte_optimization, eliminate_numeric, query_compilatio…
sfc-gh-aalam Sep 13, 2024
638dd09
cover more configs
sfc-gh-aalam Sep 17, 2024
7ae2c33
fix SnowflakePlan copy
sfc-gh-aalam Sep 17, 2024
1689ebf
minor update
sfc-gh-aalam Sep 17, 2024
5e8a2d2
add description
sfc-gh-aalam Sep 17, 2024
1c83ef2
use _package_lock to protect Session._packages
sfc-gh-aalam Sep 17, 2024
a649761
undo refactor
sfc-gh-aalam Sep 17, 2024
f03d618
undo refactor
sfc-gh-aalam Sep 17, 2024
5f398d5
fix test
sfc-gh-aalam Sep 17, 2024
3807087
fix test
sfc-gh-aalam Sep 17, 2024
4eef3e9
Merge branch 'aalam-SNOW-1418523-make-internal-session-variables-thre…
sfc-gh-aalam Sep 17, 2024
df3263c
add file IO tests
sfc-gh-aalam Sep 12, 2024
6769c54
merge with base
sfc-gh-aalam Sep 17, 2024
af86f67
merge with base
sfc-gh-aalam Sep 17, 2024
a737f33
fix test
sfc-gh-aalam Sep 17, 2024
8ca2730
protect complexity bounds setter with lock
sfc-gh-aalam Sep 17, 2024
81417a3
add config context
sfc-gh-aalam Sep 19, 2024
e340567
add tests
sfc-gh-aalam Sep 19, 2024
30952bb
update documentation
sfc-gh-aalam Sep 20, 2024
03f25b5
use config context in plan compiler
sfc-gh-aalam Sep 20, 2024
6deb402
add comments
sfc-gh-aalam Sep 20, 2024
8e1dfe0
minor refactor
sfc-gh-aalam Sep 20, 2024
10bfeb4
fix test
sfc-gh-aalam Sep 20, 2024
879940a
update documentation
sfc-gh-aalam Sep 20, 2024
5aad2d9
simplify context config
sfc-gh-aalam Sep 25, 2024
669eb91
merge with base
sfc-gh-aalam Sep 25, 2024
a85a144
add config context to repeated subquery elimination resolution stage
sfc-gh-aalam Sep 25, 2024
a79ffb4
fix tests
sfc-gh-aalam Sep 26, 2024
4420350
refactor
sfc-gh-aalam Sep 26, 2024
5f1eaa6
remove do_analyze
sfc-gh-aalam Sep 27, 2024
9d62017
fix
sfc-gh-aalam Sep 27, 2024
b58aa8b
fix
sfc-gh-aalam Sep 27, 2024
db37033
fix
sfc-gh-aalam Sep 27, 2024
dddd15f
fix unit tests
sfc-gh-aalam Sep 27, 2024
57ee9e8
simplify
sfc-gh-aalam Sep 27, 2024
809a86e
simplify
sfc-gh-aalam Sep 27, 2024
6021ab8
simplify
sfc-gh-aalam Sep 27, 2024
43986f6
simplify
sfc-gh-aalam Sep 27, 2024
0430e92
simplify
sfc-gh-aalam Sep 27, 2024
095b04e
remove config context
sfc-gh-aalam Sep 30, 2024
32707f9
min-diff
sfc-gh-aalam Sep 30, 2024
3bf678d
min-diff
sfc-gh-aalam Sep 30, 2024
3eade1a
min-diff
sfc-gh-aalam Sep 30, 2024
1850d5d
Merge branch 'aalam-SNOW-1418523-make-internal-session-variables-thre…
sfc-gh-aalam Oct 2, 2024
1fa6ad2
add warnings
sfc-gh-aalam Oct 2, 2024
7c85432
Merge branch 'aalam-SNOW-1418523-make-internal-session-variables-thre…
sfc-gh-aalam Oct 3, 2024
f994842
address feedback
sfc-gh-aalam Oct 3, 2024
4621836
address feedback
sfc-gh-aalam Oct 3, 2024
e1c68f3
fix string
sfc-gh-aalam Oct 3, 2024
e5b48dd
ignore on multi-thread
sfc-gh-aalam Oct 3, 2024
496e2be
undo ignore
sfc-gh-aalam Oct 3, 2024
980d3b7
update warning message
sfc-gh-aalam Oct 4, 2024
54a6b5d
address comments
sfc-gh-aalam Oct 4, 2024
67609e8
address comments
sfc-gh-aalam Oct 4, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
address feedback
sfc-gh-aalam committed Oct 3, 2024
commit f9948424382b107819241c4f465a5ae8d514d9e7
32 changes: 32 additions & 0 deletions src/snowflake/snowpark/session.py
Original file line number Diff line number Diff line change
@@ -781,6 +781,13 @@ def custom_package_usage_config(self) -> Dict:

@sql_simplifier_enabled.setter
def sql_simplifier_enabled(self, value: bool) -> None:
if threading.active_count() > 1:
_logger.warning(
"Setting sql_simplifier_enabled is not currently thread-safe. "
"Ignoring the update"
)
return

with self._lock:
self._conn._telemetry_client.send_sql_simplifier_telemetry(
self._session_id, value
@@ -802,6 +809,7 @@ def cte_optimization_enabled(self, value: bool) -> None:
"Setting cte_optimization_enabled is not currently thread-safe. Ignoring the update"
)
return

with self._lock:
if value:
self._conn._telemetry_client.send_cte_optimization_telemetry(
@@ -813,6 +821,12 @@ def cte_optimization_enabled(self, value: bool) -> None:
@experimental_parameter(version="1.20.0")
def eliminate_numeric_sql_value_cast_enabled(self, value: bool) -> None:
"""Set the value for eliminate_numeric_sql_value_cast_enabled"""
if threading.active_count() > 1:
_logger.warning(
"Setting eliminate_numeric_sql_value_cast_enabled is not currently thread-safe. "
"Ignoring the update"
)
return

if value in [True, False]:
with self._lock:
@@ -829,6 +843,13 @@ def eliminate_numeric_sql_value_cast_enabled(self, value: bool) -> None:
@experimental_parameter(version="1.21.0")
def auto_clean_up_temp_table_enabled(self, value: bool) -> None:
"""Set the value for auto_clean_up_temp_table_enabled"""
if threading.active_count() > 1:
_logger.warning(
"Setting auto_clean_up_temp_table_enabled is not currently thread-safe. "
"Ignoring the update"
)
return

if value in [True, False]:
self._conn._telemetry_client.send_auto_clean_up_temp_table_telemetry(
self._session_id, value
@@ -847,6 +868,11 @@ def large_query_breakdown_enabled(self, value: bool) -> None:
materialize the partitions, and then combine them to execute the query to improve
overall performance.
"""
if threading.active_count() > 1:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you can probably have a small utility function and call that for all configs with the config name as a parameter

_logger.warning(
"Setting large_query_breakdown_enabled is not currently thread-safe. Ignoring the update"
)
return

if value in [True, False]:
with self._lock:
@@ -862,6 +888,12 @@ def large_query_breakdown_enabled(self, value: bool) -> None:
@large_query_breakdown_complexity_bounds.setter
def large_query_breakdown_complexity_bounds(self, value: Tuple[int, int]) -> None:
"""Set the lower and upper bounds for the complexity score used in large query breakdown optimization."""
if threading.active_count() > 1:
_logger.warning(
"Setting large_query_breakdown_complexity_bounds is not currently thread-safe. "
"Ignoring the update"
)
return

if len(value) != 2:
raise ValueError(
34 changes: 23 additions & 11 deletions tests/integ/test_multithreading.py
Original file line number Diff line number Diff line change
@@ -505,18 +505,30 @@ def finish(self):
executor.submit(register_and_test_udaf, session, i)


def test_concurrent_update_on_cte_optimization_enabled(session, caplog):
def run_cte_optimization(session_, thread_id):
if thread_id % 2 == 0:
session_.cte_optimization_enabled = True
else:
session_.cte_optimization_enabled = False
@pytest.mark.parametrize(
"config,value",
[
("cte_optimization_enabled", True),
("sql_simplifier_enabled", True),
("eliminate_numeric_sql_value_cast_enabled", True),
("auto_clean_up_temp_table_enabled", True),
("large_query_breakdown_enabled", True),
("large_query_breakdown_complexity_bounds", (20, 30)),
],
)
def test_concurrent_update_on_sensitive_configs(session, config, value, caplog):
def change_config_value(session_):
session_.conf.set(config, value)

caplog.clear()

# check everything works find outside multiple threads
with caplog.at_level(logging.WARNING):
change_config_value(session)
assert f"Setting {config} is not currently thread-safe" not in caplog.text

with caplog.at_level(logging.WARNING):
with ThreadPoolExecutor(max_workers=5) as executor:
for i in range(5):
executor.submit(run_cte_optimization, session, i)
assert (
"Setting cte_optimization_enabled is not currently thread-safe" in caplog.text
)
for _ in range(5):
executor.submit(change_config_value, session)
assert f"Setting {config} is not currently thread-safe" in caplog.text