Skip to content

Commit 43555a7

Browse files
ptiurinptiurin
authored
fix(NoTicket): Propagating parameters from USE ENGINE (#469)
1 parent 18c93ed commit 43555a7

File tree

9 files changed

+315
-8
lines changed

9 files changed

+315
-8
lines changed

‎src/firebolt/async_db/connection.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -354,7 +354,7 @@ async def connect_v2(
354354
client.clone(),
355355
CursorV2,
356356
api_endpoint,
357-
cursor.parameters,
357+
cursor.parameters | cursor._set_parameters,
358358
connection_id,
359359
)
360360

‎src/firebolt/async_db/cursor.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -346,7 +346,9 @@ async def use_engine(self, engine: str, cache: bool = True) -> None:
346346
self._update_set_parameters(cache_obj.engines[engine].params)
347347
else:
348348
await self.execute(f'USE ENGINE "{engine}"')
349-
cache_obj.engines[engine] = EngineInfo(self.engine_url, self.parameters)
349+
cache_obj.engines[engine] = EngineInfo(
350+
self.engine_url, self.parameters | self._set_parameters
351+
)
350352
self.set_cache_record(cache_obj)
351353
else:
352354
await self.execute(f'USE ENGINE "{engine}"')

‎src/firebolt/db/connection.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ def connect_v2(
183183
client.clone(),
184184
CursorV2,
185185
api_endpoint,
186-
cursor.parameters,
186+
cursor.parameters | cursor._set_parameters,
187187
connection_id,
188188
)
189189

‎src/firebolt/db/cursor.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -350,7 +350,9 @@ def use_engine(self, engine: str, cache: bool = True) -> None:
350350
self._update_set_parameters(cache_obj.engines[engine].params)
351351
else:
352352
self.execute(f'USE ENGINE "{engine}"')
353-
cache_obj.engines[engine] = EngineInfo(self.engine_url, self.parameters)
353+
cache_obj.engines[engine] = EngineInfo(
354+
self.engine_url, self.parameters | self._set_parameters
355+
)
354356
self.set_cache_record(cache_obj)
355357
else:
356358
self.execute(f'USE ENGINE "{engine}"')

‎tests/unit/async_db/test_caching.py‎

Lines changed: 77 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import time
2-
from typing import Callable, Generator
2+
from typing import Callable, Dict, Generator
33
from unittest.mock import patch
44

55
from httpx import URL
@@ -551,3 +551,79 @@ def use_engine_callback_counter(request, **kwargs):
551551
assert (
552552
use_engine_call_counter == 2
553553
), "Use engine URL was not called after cache expiry"
554+
555+
556+
async def test_use_engine_parameters_caching(
557+
db_name: str,
558+
engine_name: str,
559+
auth_url: str,
560+
api_endpoint: str,
561+
auth: Auth,
562+
account_name: str,
563+
httpx_mock: HTTPXMock,
564+
system_engine_no_db_query_url: str,
565+
system_engine_query_url: str,
566+
use_database_callback: Callable,
567+
use_engine_with_params_callback: Callable,
568+
test_update_parameters: Dict[str, str],
569+
mock_system_engine_connection_flow: Callable,
570+
):
571+
"""Test that USE ENGINE parameters are cached and correctly retrieved."""
572+
mock_system_engine_connection_flow()
573+
574+
use_engine_call_counter = 0
575+
576+
def use_engine_callback_counter(request, **kwargs):
577+
nonlocal use_engine_call_counter
578+
use_engine_call_counter += 1
579+
return use_engine_with_params_callback(request, **kwargs)
580+
581+
# Add the missing USE DATABASE callback
582+
httpx_mock.add_callback(
583+
use_database_callback,
584+
url=system_engine_no_db_query_url,
585+
match_content=f'USE DATABASE "{db_name}"'.encode("utf-8"),
586+
is_reusable=True,
587+
)
588+
589+
# Add USE ENGINE callback with parameters
590+
httpx_mock.add_callback(
591+
use_engine_callback_counter,
592+
url=system_engine_query_url,
593+
match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"),
594+
is_reusable=True,
595+
)
596+
597+
# First connection - should populate cache with parameters
598+
async with await connect(
599+
database=db_name,
600+
engine_name=engine_name,
601+
auth=auth,
602+
account_name=account_name,
603+
api_endpoint=api_endpoint,
604+
) as connection:
605+
cursor = connection.cursor()
606+
# Verify parameters are set in cursor from USE ENGINE response
607+
for param_name, expected_value in test_update_parameters.items():
608+
assert param_name in cursor._set_parameters
609+
assert cursor._set_parameters[param_name] == expected_value
610+
611+
# Verify USE ENGINE was called once
612+
assert use_engine_call_counter == 1, "USE ENGINE was not called on first connection"
613+
614+
# Second connection - should use cache and not call USE ENGINE again
615+
async with await connect(
616+
database=db_name,
617+
engine_name=engine_name,
618+
auth=auth,
619+
account_name=account_name,
620+
api_endpoint=api_endpoint,
621+
) as connection:
622+
cursor = connection.cursor()
623+
# Verify cached parameters are correctly applied
624+
for param_name, expected_value in test_update_parameters.items():
625+
assert param_name in cursor._set_parameters
626+
assert cursor._set_parameters[param_name] == expected_value
627+
628+
# Verify USE ENGINE was not called again (cache hit)
629+
assert use_engine_call_counter == 1, "USE ENGINE was called when cache should hit"

‎tests/unit/async_db/test_connection.py‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import Callable, Generator, List, Optional, Tuple
1+
from typing import Callable, Dict, Generator, List, Optional, Tuple
22
from unittest.mock import ANY as AnyValue
33
from unittest.mock import MagicMock, patch
44

@@ -765,3 +765,55 @@ async def test_multiple_results_for_async_token(
765765
assert len(query_info) == 2, "Expected two results for the same token"
766766
assert query_info[0].query_id == async_multiple_query_data[0][7]
767767
assert query_info[1].query_id == async_multiple_query_data[1][7]
768+
769+
770+
async def test_use_engine_update_parameters_propagation(
771+
db_name: str,
772+
account_name: str,
773+
engine_name: str,
774+
auth: Auth,
775+
api_endpoint: str,
776+
httpx_mock: HTTPXMock,
777+
system_engine_no_db_query_url: str,
778+
system_engine_query_url: str,
779+
use_database_callback: Callable,
780+
use_engine_with_params_callback: Callable,
781+
test_update_parameters: Dict[str, str],
782+
mock_system_engine_connection_flow: Callable,
783+
) -> None:
784+
"""Test that USE ENGINE with Firebolt-Update-Parameters header propagates to connection init_parameters and cursors."""
785+
mock_system_engine_connection_flow()
786+
787+
# Mock USE DATABASE callback
788+
httpx_mock.add_callback(
789+
use_database_callback,
790+
url=system_engine_no_db_query_url,
791+
match_content=f'USE DATABASE "{db_name}"'.encode("utf-8"),
792+
is_reusable=True,
793+
)
794+
795+
# Mock USE ENGINE callback with parameter updates
796+
httpx_mock.add_callback(
797+
use_engine_with_params_callback,
798+
url=system_engine_query_url,
799+
match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"),
800+
is_reusable=True,
801+
)
802+
803+
async with await connect(
804+
database=db_name,
805+
auth=auth,
806+
engine_name=engine_name,
807+
account_name=account_name,
808+
api_endpoint=api_endpoint,
809+
) as connection:
810+
# Verify that parameters from USE ENGINE header are in connection init_parameters
811+
for param_name, expected_value in test_update_parameters.items():
812+
assert param_name in connection.init_parameters
813+
assert connection.init_parameters[param_name] == expected_value
814+
815+
# Verify that new cursors get these parameters
816+
cursor = connection.cursor()
817+
for param_name, expected_value in test_update_parameters.items():
818+
assert param_name in cursor._set_parameters
819+
assert cursor._set_parameters[param_name] == expected_value

‎tests/unit/db/test_caching.py‎

Lines changed: 77 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import time
2-
from typing import Callable, Generator
2+
from typing import Callable, Dict, Generator
33
from unittest.mock import patch
44

55
from httpx import URL
@@ -551,3 +551,79 @@ def use_engine_callback_counter(request, **kwargs):
551551
assert (
552552
use_engine_call_counter == 2
553553
), "Use engine URL was not called after cache expiry"
554+
555+
556+
def test_use_engine_parameters_caching(
557+
db_name: str,
558+
engine_name: str,
559+
auth_url: str,
560+
api_endpoint: str,
561+
auth: Auth,
562+
account_name: str,
563+
httpx_mock: HTTPXMock,
564+
system_engine_no_db_query_url: str,
565+
system_engine_query_url: str,
566+
use_database_callback: Callable,
567+
use_engine_with_params_callback: Callable,
568+
test_update_parameters: Dict[str, str],
569+
mock_system_engine_connection_flow: Callable,
570+
):
571+
"""Test that USE ENGINE parameters are cached and correctly retrieved."""
572+
mock_system_engine_connection_flow()
573+
574+
use_engine_call_counter = 0
575+
576+
def use_engine_callback_counter(request, **kwargs):
577+
nonlocal use_engine_call_counter
578+
use_engine_call_counter += 1
579+
return use_engine_with_params_callback(request, **kwargs)
580+
581+
# Add the missing USE DATABASE callback
582+
httpx_mock.add_callback(
583+
use_database_callback,
584+
url=system_engine_no_db_query_url,
585+
match_content=f'USE DATABASE "{db_name}"'.encode("utf-8"),
586+
is_reusable=True,
587+
)
588+
589+
# Add USE ENGINE callback with parameters
590+
httpx_mock.add_callback(
591+
use_engine_callback_counter,
592+
url=system_engine_query_url,
593+
match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"),
594+
is_reusable=True,
595+
)
596+
597+
# First connection - should populate cache with parameters
598+
with connect(
599+
database=db_name,
600+
engine_name=engine_name,
601+
auth=auth,
602+
account_name=account_name,
603+
api_endpoint=api_endpoint,
604+
) as connection:
605+
cursor = connection.cursor()
606+
# Verify parameters are set in cursor from USE ENGINE response
607+
for param_name, expected_value in test_update_parameters.items():
608+
assert param_name in cursor._set_parameters
609+
assert cursor._set_parameters[param_name] == expected_value
610+
611+
# Verify USE ENGINE was called once
612+
assert use_engine_call_counter == 1, "USE ENGINE was not called on first connection"
613+
614+
# Second connection - should use cache and not call USE ENGINE again
615+
with connect(
616+
database=db_name,
617+
engine_name=engine_name,
618+
auth=auth,
619+
account_name=account_name,
620+
api_endpoint=api_endpoint,
621+
) as connection:
622+
cursor = connection.cursor()
623+
# Verify cached parameters are correctly applied
624+
for param_name, expected_value in test_update_parameters.items():
625+
assert param_name in cursor._set_parameters
626+
assert cursor._set_parameters[param_name] == expected_value
627+
628+
# Verify USE ENGINE was not called again (cache hit)
629+
assert use_engine_call_counter == 1, "USE ENGINE was called when cache should hit"

‎tests/unit/db/test_connection.py‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import gc
22
import warnings
3-
from typing import Callable, Generator, List, Optional, Tuple
3+
from typing import Callable, Dict, Generator, List, Optional, Tuple
44
from unittest.mock import ANY as AnyValue
55
from unittest.mock import MagicMock, patch
66

@@ -783,3 +783,55 @@ def test_multiple_results_for_async_token(
783783
assert len(query_info) == 2, "Expected two results for the same token"
784784
assert query_info[0].query_id == async_multiple_query_data[0][7]
785785
assert query_info[1].query_id == async_multiple_query_data[1][7]
786+
787+
788+
def test_use_engine_update_parameters_propagation(
789+
db_name: str,
790+
account_name: str,
791+
engine_name: str,
792+
auth: Auth,
793+
api_endpoint: str,
794+
httpx_mock: HTTPXMock,
795+
system_engine_no_db_query_url: str,
796+
system_engine_query_url: str,
797+
use_database_callback: Callable,
798+
use_engine_with_params_callback: Callable,
799+
test_update_parameters: Dict[str, str],
800+
mock_system_engine_connection_flow: Callable,
801+
) -> None:
802+
"""Test that USE ENGINE with Firebolt-Update-Parameters header propagates to connection init_parameters and cursors."""
803+
mock_system_engine_connection_flow()
804+
805+
# Mock USE DATABASE callback
806+
httpx_mock.add_callback(
807+
use_database_callback,
808+
url=system_engine_no_db_query_url,
809+
match_content=f'USE DATABASE "{db_name}"'.encode("utf-8"),
810+
is_reusable=True,
811+
)
812+
813+
# Mock USE ENGINE callback with parameter updates
814+
httpx_mock.add_callback(
815+
use_engine_with_params_callback,
816+
url=system_engine_query_url,
817+
match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"),
818+
is_reusable=True,
819+
)
820+
821+
with connect(
822+
database=db_name,
823+
auth=auth,
824+
engine_name=engine_name,
825+
account_name=account_name,
826+
api_endpoint=api_endpoint,
827+
) as connection:
828+
# Verify that parameters from USE ENGINE header are in connection init_parameters
829+
for param_name, expected_value in test_update_parameters.items():
830+
assert param_name in connection.init_parameters
831+
assert connection.init_parameters[param_name] == expected_value
832+
833+
# Verify that new cursors get these parameters
834+
cursor = connection.cursor()
835+
for param_name, expected_value in test_update_parameters.items():
836+
assert param_name in cursor._set_parameters
837+
assert cursor._set_parameters[param_name] == expected_value

0 commit comments

Comments
 (0)