fix: Better invalid account error (#331)

Author: ptiurin

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}: "

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"]

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}: "

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.
 

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)

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)

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(

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:

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)

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(