diff --git a/src/firebolt/async_db/util.py b/src/firebolt/async_db/util.py index e48ca9949a..2373a49ae3 100644 --- a/src/firebolt/async_db/util.py +++ b/src/firebolt/async_db/util.py @@ -5,7 +5,10 @@ from firebolt.client.auth import Auth from firebolt.client.client import AsyncClientV2 from firebolt.common.settings import DEFAULT_TIMEOUT_SECONDS -from firebolt.utils.exception import AccountNotFoundError, InterfaceError +from firebolt.utils.exception import ( + AccountNotFoundOrNoAccessError, + InterfaceError, +) from firebolt.utils.urls import GATEWAY_HOST_BY_ACCOUNT_NAME ENGINE_STATUS_RUNNING = "Running" @@ -26,7 +29,7 @@ async def _get_system_engine_url( url = GATEWAY_HOST_BY_ACCOUNT_NAME.format(account_name=account_name) response = await client.get(url=url) if response.status_code == codes.NOT_FOUND: - raise AccountNotFoundError(account_name) + raise AccountNotFoundOrNoAccessError(account_name) if response.status_code != codes.OK: raise InterfaceError( f"Unable to retrieve system engine endpoint {url}: " diff --git a/src/firebolt/client/client.py b/src/firebolt/client/client.py index 74c8746802..a99ecbf2a5 100644 --- a/src/firebolt/client/client.py +++ b/src/firebolt/client/client.py @@ -18,6 +18,7 @@ ) from firebolt.utils.exception import ( AccountNotFoundError, + AccountNotFoundOrNoAccessError, FireboltEngineError, InterfaceError, ) @@ -142,7 +143,7 @@ def account_id(self) -> str: ) if response.status_code == HttpxCodes.NOT_FOUND: assert self.account_name is not None - raise AccountNotFoundError(self.account_name) + raise AccountNotFoundOrNoAccessError(self.account_name) # process all other status codes response.raise_for_status() return response.json()["id"] @@ -324,7 +325,7 @@ async def account_id(self) -> str: ) if response.status_code == HttpxCodes.NOT_FOUND: assert self.account_name is not None - raise AccountNotFoundError(self.account_name) + raise AccountNotFoundOrNoAccessError(self.account_name) # process all other status codes response.raise_for_status() account_id = response.json()["id"] diff --git a/src/firebolt/db/util.py b/src/firebolt/db/util.py index c3d9676609..2bf0be27e2 100644 --- a/src/firebolt/db/util.py +++ b/src/firebolt/db/util.py @@ -5,7 +5,10 @@ from firebolt.client import ClientV2 from firebolt.client.auth import Auth from firebolt.common.settings import DEFAULT_TIMEOUT_SECONDS -from firebolt.utils.exception import AccountNotFoundError, InterfaceError +from firebolt.utils.exception import ( + AccountNotFoundOrNoAccessError, + InterfaceError, +) from firebolt.utils.urls import GATEWAY_HOST_BY_ACCOUNT_NAME ENGINE_STATUS_RUNNING = "Running" @@ -26,7 +29,7 @@ def _get_system_engine_url( url = GATEWAY_HOST_BY_ACCOUNT_NAME.format(account_name=account_name) response = client.get(url=url) if response.status_code == codes.NOT_FOUND: - raise AccountNotFoundError(account_name) + raise AccountNotFoundOrNoAccessError(account_name) if response.status_code != codes.OK: raise InterfaceError( f"Unable to retrieve system engine endpoint {url}: " diff --git a/src/firebolt/utils/exception.py b/src/firebolt/utils/exception.py index fee1d9f104..bb9877941a 100644 --- a/src/firebolt/utils/exception.py +++ b/src/firebolt/utils/exception.py @@ -75,6 +75,27 @@ def __init__(self, account_name: str): self.account_name = account_name +class AccountNotFoundOrNoAccessError(FireboltError): + """Account with provided name doesn't exist. + + Args: + account_name (str): Name of account that wasn't found + + Attributes: + account_name (str): Name of account that wasn't found + """ + + def __init__(self, account_name: str): + super().__init__( + f"Account '{account_name}' does not exist " + "in this organization or is not authorized. " + "Please verify the account name and make sure your " + "service account has the correct RBAC permissions and " + "is linked to a user." + ) + self.account_name = account_name + + class AttachedEngineInUseError(FireboltDatabaseError): """Engine unavailable because it's starting/stopping. diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 8b612a1de1..f1766cbb08 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -94,6 +94,11 @@ def account_name() -> str: return must_env(ACCOUNT_NAME_ENV) +@fixture(scope="session") +def invalid_account_name(account_name: str) -> str: + return f"{account_name}--" + + @fixture(scope="session") def api_endpoint() -> str: return must_env(API_ENDPOINT_ENV) diff --git a/tests/integration/dbapi/async/V2/conftest.py b/tests/integration/dbapi/async/V2/conftest.py index 3bfbbc7ff4..a32f1e1758 100644 --- a/tests/integration/dbapi/async/V2/conftest.py +++ b/tests/integration/dbapi/async/V2/conftest.py @@ -1,7 +1,12 @@ +from random import randint +from typing import Tuple + from pytest import fixture from firebolt.async_db import Connection, connect from firebolt.client.auth.base import Auth +from firebolt.client.auth.client_credentials import ClientCredentials +from tests.integration.conftest import Secret @fixture @@ -66,3 +71,38 @@ async def connection_system_engine_no_db( api_endpoint=api_endpoint, ) as connection: yield connection + + +@fixture +async def service_account_no_user( + connection_system_engine_no_db: Connection, + database_name: str, +) -> Tuple[str, Secret]: + # function-level fixture so we need to make sa name is unique + sa_account_name = f"{database_name}_sa_no_user_{randint(0, 1000)}" + with connection_system_engine_no_db.cursor() as cursor: + await cursor.execute( + f'CREATE SERVICE ACCOUNT "{sa_account_name}" ' + 'WITH DESCRIPTION = "Ecosytem test with no user"' + ) + await cursor.execute(f"CALL fb_GENERATESERVICEACCOUNTKEY('{sa_account_name}')") + # service_account_name, service_account_id, secret + _, s_id, key = await cursor.fetchone() + # Currently this is bugged so retrieve id via a query. FIR-28719 + if not s_id: + await cursor.execute( + "SELECT service_account_id FROM information_schema.service_accounts " + f"WHERE service_account_name='{sa_account_name}'" + ) + s_id = (await cursor.fetchone())[0] + # Wrap in secret to avoid leaking the key in the logs + yield s_id, Secret(key) + await cursor.execute(f"DROP SERVICE ACCOUNT {sa_account_name}") + + +@fixture +async def auth_no_user( + service_account_no_user: Tuple[str, Secret], +) -> ClientCredentials: + s_id, s_secret = service_account_no_user + return ClientCredentials(s_id, s_secret.value) diff --git a/tests/integration/dbapi/async/V2/test_errors_async.py b/tests/integration/dbapi/async/V2/test_errors_async.py index 6a569f6683..d3a31e0471 100644 --- a/tests/integration/dbapi/async/V2/test_errors_async.py +++ b/tests/integration/dbapi/async/V2/test_errors_async.py @@ -3,7 +3,7 @@ from firebolt.async_db import Connection, connect from firebolt.client.auth import ClientCredentials from firebolt.utils.exception import ( - AccountNotFoundError, + AccountNotFoundOrNoAccessError, EngineNotRunningError, FireboltEngineError, InterfaceError, @@ -13,25 +13,45 @@ async def test_invalid_account( database_name: str, - engine_name: str, + invalid_account_name: str, auth: ClientCredentials, api_endpoint: str, ) -> None: """Connection properly reacts to invalid account error.""" - account_name = "--" - with raises(AccountNotFoundError) as exc_info: + with raises(AccountNotFoundOrNoAccessError) as exc_info: async with await connect( database=database_name, - engine_name=engine_name, auth=auth, + account_name=invalid_account_name, + api_endpoint=api_endpoint, + ) as connection: + await connection.cursor().execute("show tables") + + assert str(exc_info.value).startswith( + f"Account '{invalid_account_name}' does not exist" + ), "Invalid account error message." + + +async def test_account_no_user( + database_name: str, + account_name: str, + auth_no_user: ClientCredentials, + api_endpoint: str, +) -> None: + """Connection properly reacts to account that doesn't have + a user attached to it.""" + with raises(AccountNotFoundOrNoAccessError) as exc_info: + async with await connect( + database=database_name, + auth=auth_no_user, account_name=account_name, api_endpoint=api_endpoint, ) as connection: await connection.cursor().execute("show tables") - assert str(exc_info.value).startswith( - f'Account "{account_name}" does not exist' - ), "Invalid account error message." + assert str(exc_info.value).startswith( + f"Account '{account_name}' does not exist" + ), "Invalid account error message." async def test_engine_name_not_exists( diff --git a/tests/integration/dbapi/async/V2/test_queries_async.py b/tests/integration/dbapi/async/V2/test_queries_async.py index cddc0cf951..102d5ebd42 100644 --- a/tests/integration/dbapi/async/V2/test_queries_async.py +++ b/tests/integration/dbapi/async/V2/test_queries_async.py @@ -415,7 +415,7 @@ async def test_bytea_roundtrip( ), "Invalid bytea data returned after roundtrip" -@fixture +@fixture(scope="session") async def setup_db(connection_system_engine_no_db: Connection, use_db_name: str): use_db_name = use_db_name + "_async" with connection_system_engine_no_db.cursor() as cursor: diff --git a/tests/integration/dbapi/sync/V2/conftest.py b/tests/integration/dbapi/sync/V2/conftest.py index 2b561c7fb2..be1ec69543 100644 --- a/tests/integration/dbapi/sync/V2/conftest.py +++ b/tests/integration/dbapi/sync/V2/conftest.py @@ -1,7 +1,12 @@ +from random import randint +from typing import Tuple + from pytest import fixture from firebolt.client.auth.base import Auth +from firebolt.client.auth.client_credentials import ClientCredentials from firebolt.db import Connection, connect +from tests.integration.conftest import Secret @fixture @@ -66,3 +71,38 @@ def connection_system_engine_no_db( api_endpoint=api_endpoint, ) as connection: yield connection + + +@fixture +def service_account_no_user( + connection_system_engine_no_db: Connection, + database_name: str, +) -> Tuple[str, Secret]: + # function-level fixture so we need to make sa name is unique + sa_account_name = f"{database_name}_sa_no_user_{randint(0, 1000)}" + with connection_system_engine_no_db.cursor() as cursor: + cursor.execute( + f'CREATE SERVICE ACCOUNT "{sa_account_name}" ' + 'WITH DESCRIPTION = "Ecosytem test with no user"' + ) + cursor.execute(f"CALL fb_GENERATESERVICEACCOUNTKEY('{sa_account_name}')") + # service_account_name, service_account_id, secret + _, s_id, key = cursor.fetchone() + # Currently this is bugged so retrieve id via a query. FIR-28719 + if not s_id: + cursor.execute( + "SELECT service_account_id FROM information_schema.service_accounts " + f"WHERE service_account_name='{sa_account_name}'" + ) + s_id = cursor.fetchone()[0] + # Wrap in secret to avoid leaking the key in the logs + yield s_id, Secret(key) + cursor.execute(f"DROP SERVICE ACCOUNT {sa_account_name}") + + +@fixture +def auth_no_user( + service_account_no_user: Tuple[str, Secret], +) -> ClientCredentials: + s_id, s_secret = service_account_no_user + return ClientCredentials(s_id, s_secret.value) diff --git a/tests/integration/dbapi/sync/V2/test_errors.py b/tests/integration/dbapi/sync/V2/test_errors.py index aab2621649..0080fedd69 100644 --- a/tests/integration/dbapi/sync/V2/test_errors.py +++ b/tests/integration/dbapi/sync/V2/test_errors.py @@ -3,7 +3,7 @@ from firebolt.client.auth import ClientCredentials from firebolt.db import Connection, connect from firebolt.utils.exception import ( - AccountNotFoundError, + AccountNotFoundOrNoAccessError, EngineNotRunningError, FireboltDatabaseError, FireboltEngineError, @@ -13,25 +13,44 @@ def test_invalid_account( database_name: str, - engine_name: str, + invalid_account_name: str, auth: ClientCredentials, api_endpoint: str, ) -> None: """Connection properly reacts to invalid account error.""" - account_name = "--" - with raises(AccountNotFoundError) as exc_info: + with raises(AccountNotFoundOrNoAccessError) as exc_info: with connect( database=database_name, - engine_name=engine_name, # Omit engine_url to force account_id lookup. auth=auth, + account_name=invalid_account_name, + api_endpoint=api_endpoint, + ) as connection: + connection.cursor().execute("show tables") + + assert str(exc_info.value).startswith( + f"Account '{invalid_account_name}' does not exist" + ), "Invalid account error message." + + +def test_account_no_user( + database_name: str, + account_name: str, + auth_no_user: ClientCredentials, + api_endpoint: str, +) -> None: + """Connection properly reacts to invalid account error.""" + with raises(AccountNotFoundOrNoAccessError) as exc_info: + with connect( + database=database_name, + auth=auth_no_user, account_name=account_name, api_endpoint=api_endpoint, ) as connection: connection.cursor().execute("show tables") - assert str(exc_info.value).startswith( - f'Account "{account_name}" does not exist' - ), "Invalid account error message." + assert str(exc_info.value).startswith( + f"Account '{account_name}' does not exist" + ), "Invalid account error message." def test_engine_name_not_exists( diff --git a/tests/unit/async_db/V2/test_connection.py b/tests/unit/async_db/V2/test_connection.py index d2a8dfaa4d..2fa1745922 100644 --- a/tests/unit/async_db/V2/test_connection.py +++ b/tests/unit/async_db/V2/test_connection.py @@ -9,6 +9,7 @@ from firebolt.client.auth import Auth, ClientCredentials from firebolt.common._types import ColType from firebolt.utils.exception import ( + AccountNotFoundOrNoAccessError, ConfigurationError, ConnectionClosedError, EngineNotRunningError, @@ -178,6 +179,55 @@ async def test_connect_database( assert await connection.cursor().execute("select*") == len(python_query_data) +async def test_connect_invalid_account( + db_name: str, + auth_url: str, + server: str, + auth: Auth, + account_name: str, + httpx_mock: HTTPXMock, + check_credentials_callback: Callable, + get_system_engine_url: str, + get_system_engine_callback: Callable, + account_id_url: str, + account_id_invalid_callback: Callable, +): + httpx_mock.add_callback(check_credentials_callback, url=auth_url) + httpx_mock.add_callback(get_system_engine_callback, url=get_system_engine_url) + httpx_mock.add_callback(account_id_invalid_callback, url=account_id_url) + with raises(AccountNotFoundOrNoAccessError): + async with await connect( + database=db_name, + auth=auth, + account_name=account_name, + api_endpoint=server, + ) as connection: + await connection.cursor().execute("select*") + + +async def test_connect_system_engine_404( + db_name: str, + auth_url: str, + server: str, + auth: Auth, + account_name: str, + httpx_mock: HTTPXMock, + check_credentials_callback: Callable, + get_system_engine_url: str, + get_system_engine_404_callback: Callable, +): + httpx_mock.add_callback(check_credentials_callback, url=auth_url) + httpx_mock.add_callback(get_system_engine_404_callback, url=get_system_engine_url) + with raises(AccountNotFoundOrNoAccessError): + async with await connect( + database=db_name, + auth=auth, + account_name=account_name, + api_endpoint=server, + ) as connection: + await connection.cursor().execute("select*") + + async def test_connection_commit(connection: Connection): # nothing happens connection.commit() diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index aaf9503983..a4f6d9c304 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -176,6 +176,17 @@ def do_mock( return do_mock +@fixture +def account_id_invalid_callback() -> Callable: + def do_mock( + request: Request, + **kwargs, + ) -> Response: + return Response(status_code=httpx.codes.NOT_FOUND, json={"error": "not found"}) + + return do_mock + + @fixture def engine_id() -> str: return "mock_engine_id" diff --git a/tests/unit/db/V2/test_connection.py b/tests/unit/db/V2/test_connection.py index 494b2a6923..5732726221 100644 --- a/tests/unit/db/V2/test_connection.py +++ b/tests/unit/db/V2/test_connection.py @@ -13,6 +13,7 @@ from firebolt.db import Connection, connect from firebolt.db.cursor import CursorV2 from firebolt.utils.exception import ( + AccountNotFoundOrNoAccessError, ConfigurationError, ConnectionClosedError, EngineNotRunningError, @@ -183,6 +184,55 @@ def test_connect_database( assert connection.cursor().execute("select*") == len(python_query_data) +def test_connect_invalid_account( + db_name: str, + auth_url: str, + server: str, + auth: Auth, + account_name: str, + httpx_mock: HTTPXMock, + check_credentials_callback: Callable, + get_system_engine_url: str, + get_system_engine_callback: Callable, + account_id_url: str, + account_id_invalid_callback: Callable, +): + httpx_mock.add_callback(check_credentials_callback, url=auth_url) + httpx_mock.add_callback(get_system_engine_callback, url=get_system_engine_url) + httpx_mock.add_callback(account_id_invalid_callback, url=account_id_url) + with raises(AccountNotFoundOrNoAccessError): + with connect( + database=db_name, + auth=auth, + account_name=account_name, + api_endpoint=server, + ) as connection: + connection.cursor().execute("select*") + + +def test_connect_system_engine_404( + db_name: str, + auth_url: str, + server: str, + auth: Auth, + account_name: str, + httpx_mock: HTTPXMock, + check_credentials_callback: Callable, + get_system_engine_url: str, + get_system_engine_404_callback: Callable, +): + httpx_mock.add_callback(check_credentials_callback, url=auth_url) + httpx_mock.add_callback(get_system_engine_404_callback, url=get_system_engine_url) + with raises(AccountNotFoundOrNoAccessError): + with connect( + database=db_name, + auth=auth, + account_name=account_name, + api_endpoint=server, + ) as connection: + connection.cursor().execute("select*") + + def test_connection_unclosed_warnings(auth: Auth): c = Connection("", "", auth, "", None) with warns(UserWarning) as winfo: diff --git a/tests/unit/db_conftest.py b/tests/unit/db_conftest.py index 9a1c7a15b2..3fd04a2ed8 100644 --- a/tests/unit/db_conftest.py +++ b/tests/unit/db_conftest.py @@ -502,6 +502,22 @@ def inner( return inner +@fixture +def get_system_engine_404_callback() -> Callable: + def inner( + request: Request = None, + **kwargs, + ) -> Response: + assert request.method == "GET", "invalid request method" + + return Response( + status_code=codes.NOT_FOUND, + json={"error": "not found"}, + ) + + return inner + + @fixture def mock_connection_flow( httpx_mock: HTTPXMock,