From 121b0139c219d2218215ebc951cb9a4ce4a71e28 Mon Sep 17 00:00:00 2001 From: ptiurin Date: Wed, 7 Feb 2024 15:14:07 +0000 Subject: [PATCH 1/2] feat: Disallow internal set parameters --- src/firebolt/async_db/cursor.py | 2 ++ src/firebolt/common/base_cursor.py | 24 ++++++++++++++++++++++++ src/firebolt/db/cursor.py | 2 ++ tests/unit/async_db/V1/test_cursor.py | 22 +++++++++++++++++++++- tests/unit/async_db/V2/test_cursor.py | 22 +++++++++++++++++++++- tests/unit/db/V1/test_cursor.py | 22 +++++++++++++++++++++- tests/unit/db/V2/test_cursor.py | 22 +++++++++++++++++++++- 7 files changed, 112 insertions(+), 4 deletions(-) diff --git a/src/firebolt/async_db/cursor.py b/src/firebolt/async_db/cursor.py index 66d1579e89..04ff37078e 100644 --- a/src/firebolt/async_db/cursor.py +++ b/src/firebolt/async_db/cursor.py @@ -37,6 +37,7 @@ CursorState, QueryStatus, Statistics, + _raise_if_internal_set_parameter, check_not_closed, check_query_executed, ) @@ -126,6 +127,7 @@ async def _raise_if_error(self, resp: Response) -> None: async def _validate_set_parameter(self, parameter: SetParameter) -> None: """Validate parameter by executing simple query with it.""" + _raise_if_internal_set_parameter(parameter) if parameter.name == "async_execution": raise AsyncExecutionUnavailableError( "It is not possible to set async_execution using a SET command. " diff --git a/src/firebolt/common/base_cursor.py b/src/firebolt/common/base_cursor.py index 57a451f824..e52287a766 100644 --- a/src/firebolt/common/base_cursor.py +++ b/src/firebolt/common/base_cursor.py @@ -20,6 +20,7 @@ ) from firebolt.utils.exception import ( AsyncExecutionUnavailableError, + ConfigurationError, CursorClosedError, DataError, QueryNotRunError, @@ -55,6 +56,29 @@ class QueryStatus(Enum): # known parameters that can be set on the server side SERVER_SIDE_PARAMETERS = ["database"] +# Parameters that should be set using USE instead of SET +USE_PARAMETER_LIST = ["database", "engine"] +# parameters that can only be set by the backend +DISALLOWED_PARAMETER_LIST = ["account_id", "output_format"] + + +def _raise_if_internal_set_parameter(parameter: SetParameter) -> None: + """ + Check if parameter is internal and raise an error if it is. + """ + if parameter.name in USE_PARAMETER_LIST: + raise ConfigurationError( + "Could not set parameter. " + f"Set parameter '{parameter.name}' is not allowed. " + "Try again with 'USE ' instead of SET" + ) + if parameter.name in DISALLOWED_PARAMETER_LIST: + raise ConfigurationError( + "Could not set parameter. " + f"Set parameter '{parameter.name}' is not allowed. " + "Try again with a different parameter name." + ) + @dataclass class Statistics: diff --git a/src/firebolt/db/cursor.py b/src/firebolt/db/cursor.py index ddf2bf48c6..6117455b5a 100644 --- a/src/firebolt/db/cursor.py +++ b/src/firebolt/db/cursor.py @@ -32,6 +32,7 @@ CursorState, QueryStatus, Statistics, + _raise_if_internal_set_parameter, check_not_closed, check_query_executed, ) @@ -123,6 +124,7 @@ def _api_request( def _validate_set_parameter(self, parameter: SetParameter) -> None: """Validate parameter by executing simple query with it.""" + _raise_if_internal_set_parameter(parameter) if parameter.name == "async_execution": raise AsyncExecutionUnavailableError( "It is not possible to set async_execution using a SET command. " diff --git a/tests/unit/async_db/V1/test_cursor.py b/tests/unit/async_db/V1/test_cursor.py index e0872434a9..687eb6470f 100644 --- a/tests/unit/async_db/V1/test_cursor.py +++ b/tests/unit/async_db/V1/test_cursor.py @@ -2,7 +2,7 @@ from unittest.mock import patch from httpx import HTTPStatusError, StreamError, codes -from pytest import LogCaptureFixture, raises +from pytest import LogCaptureFixture, mark, raises from pytest_httpx import HTTPXMock from firebolt.async_db import Cursor @@ -10,6 +10,7 @@ from firebolt.common.base_cursor import ColType, CursorState, QueryStatus from firebolt.utils.exception import ( AsyncExecutionUnavailableError, + ConfigurationError, CursorClosedError, DataError, EngineNotRunningError, @@ -878,3 +879,22 @@ async def test_cursor_unknown_error_body_logging( with raises(HTTPStatusError): await cursor.execute("select 1") assert actual_error_body in caplog.text + + +@mark.parametrize( + "parameter", + [ + "database", + "engine", + "account_id", + "output_format", + ], +) +async def test_disallowed_set_parameter(cursor: Cursor, parameter: str) -> None: + """Test that setting disallowed parameters raises an error.""" + with raises(ConfigurationError) as e: + await cursor.execute(f"SET {parameter}=dummy") + assert f"Set parameter '{parameter}' is not allowed" in str( + e.value + ), "invalid error" + assert cursor._set_parameters == {}, "set parameters should not be updated" diff --git a/tests/unit/async_db/V2/test_cursor.py b/tests/unit/async_db/V2/test_cursor.py index e4ce73ac57..f0062c65ba 100644 --- a/tests/unit/async_db/V2/test_cursor.py +++ b/tests/unit/async_db/V2/test_cursor.py @@ -2,7 +2,7 @@ from unittest.mock import patch from httpx import HTTPStatusError, StreamError, codes -from pytest import LogCaptureFixture, raises +from pytest import LogCaptureFixture, mark, raises from pytest_httpx import HTTPXMock from firebolt.async_db import Cursor @@ -10,6 +10,7 @@ from firebolt.common.base_cursor import ColType, CursorState, QueryStatus from firebolt.utils.exception import ( AsyncExecutionUnavailableError, + ConfigurationError, CursorClosedError, DataError, EngineNotRunningError, @@ -893,3 +894,22 @@ async def test_cursor_unknown_error_body_logging( with raises(HTTPStatusError): await cursor.execute("select 1") assert actual_error_body in caplog.text + + +@mark.parametrize( + "parameter", + [ + "database", + "engine", + "account_id", + "output_format", + ], +) +async def test_disallowed_set_parameter(cursor: Cursor, parameter: str) -> None: + """Test that setting disallowed parameters raises an error.""" + with raises(ConfigurationError) as e: + await cursor.execute(f"SET {parameter}=dummy") + assert f"Set parameter '{parameter}' is not allowed" in str( + e.value + ), "invalid error" + assert cursor._set_parameters == {}, "set parameters should not be updated" diff --git a/tests/unit/db/V1/test_cursor.py b/tests/unit/db/V1/test_cursor.py index 66aee49cc7..009e649404 100644 --- a/tests/unit/db/V1/test_cursor.py +++ b/tests/unit/db/V1/test_cursor.py @@ -2,12 +2,13 @@ from unittest.mock import patch from httpx import HTTPStatusError, StreamError, codes -from pytest import LogCaptureFixture, raises +from pytest import LogCaptureFixture, mark, raises from pytest_httpx import HTTPXMock from firebolt.db import Cursor from firebolt.db.cursor import ColType, Column, CursorState, QueryStatus from firebolt.utils.exception import ( + ConfigurationError, CursorClosedError, DataError, OperationalError, @@ -737,3 +738,22 @@ def test_cursor_unknown_error_body_logging( with raises(HTTPStatusError): cursor.execute("select 1") assert actual_error_body in caplog.text + + +@mark.parametrize( + "parameter", + [ + "database", + "engine", + "account_id", + "output_format", + ], +) +def test_disallowed_set_parameter(cursor: Cursor, parameter: str) -> None: + """Test that setting disallowed parameters raises an error.""" + with raises(ConfigurationError) as e: + cursor.execute(f"SET {parameter}=dummy") + assert f"Set parameter '{parameter}' is not allowed" in str( + e.value + ), "invalid error" + assert cursor._set_parameters == {}, "set parameters should not be updated" diff --git a/tests/unit/db/V2/test_cursor.py b/tests/unit/db/V2/test_cursor.py index ab4f755db7..df2b091369 100644 --- a/tests/unit/db/V2/test_cursor.py +++ b/tests/unit/db/V2/test_cursor.py @@ -2,12 +2,13 @@ from unittest.mock import patch from httpx import HTTPStatusError, StreamError, codes -from pytest import LogCaptureFixture, raises +from pytest import LogCaptureFixture, mark, raises from pytest_httpx import HTTPXMock from firebolt.db import Cursor from firebolt.db.cursor import ColType, Column, CursorState, QueryStatus from firebolt.utils.exception import ( + ConfigurationError, CursorClosedError, DataError, EngineNotRunningError, @@ -798,3 +799,22 @@ def test_cursor_unknown_error_body_logging( with raises(HTTPStatusError): cursor.execute("select 1") assert actual_error_body in caplog.text + + +@mark.parametrize( + "parameter", + [ + "database", + "engine", + "account_id", + "output_format", + ], +) +def test_disallowed_set_parameter(cursor: Cursor, parameter: str) -> None: + """Test that setting disallowed parameters raises an error.""" + with raises(ConfigurationError) as e: + cursor.execute(f"SET {parameter}=dummy") + assert f"Set parameter '{parameter}' is not allowed" in str( + e.value + ), "invalid error" + assert cursor._set_parameters == {}, "set parameters should not be updated" From 4736e8e80235e585a7026d779d77f4b2bc4965f6 Mon Sep 17 00:00:00 2001 From: ptiurin Date: Thu, 8 Feb 2024 11:28:36 +0000 Subject: [PATCH 2/2] better error message --- src/firebolt/common/base_cursor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/firebolt/common/base_cursor.py b/src/firebolt/common/base_cursor.py index e52287a766..ea40b869dc 100644 --- a/src/firebolt/common/base_cursor.py +++ b/src/firebolt/common/base_cursor.py @@ -70,7 +70,7 @@ def _raise_if_internal_set_parameter(parameter: SetParameter) -> None: raise ConfigurationError( "Could not set parameter. " f"Set parameter '{parameter.name}' is not allowed. " - "Try again with 'USE ' instead of SET" + f"Try again with 'USE {str(parameter.name).upper()}' instead of SET" ) if parameter.name in DISALLOWED_PARAMETER_LIST: raise ConfigurationError(