feat: Fir 10511 feature request support resour (#128)

* add auth_token settings field

* add auth.from_token method

* allow passing access_token to resource manager

* allow connecting with access token

* add tests for access_token connection

* extend unit tests

* add ConfigurationError

Author: stepansergeevitch

src/firebolt/async_db/connection.py

@@ -8,10 +8,12 @@
 from httpcore.backends.auto import AutoBackend
 from httpcore.backends.base import AsyncNetworkStream
 from httpx import AsyncHTTPTransport, HTTPStatusError, RequestError, Timeout
+from httpx._types import AuthTypes
 
 from firebolt.async_db.cursor import BaseCursor, Cursor
-from firebolt.client import DEFAULT_API_URL, AsyncClient
+from firebolt.client import DEFAULT_API_URL, AsyncClient, Auth
 from firebolt.common.exception import (
+    ConfigurationError,
     ConnectionClosedError,
     FireboltEngineError,
     InterfaceError,
@@ -26,13 +28,12 @@
 
 async def _resolve_engine_url(
     engine_name: str,
-    username: str,
-    password: str,
+    auth: AuthTypes,
     api_endpoint: str,
     account_name: Optional[str] = None,
 ) -> str:
     async with AsyncClient(
-        auth=(username, password),
+        auth=auth,
         base_url=api_endpoint,
         account_name=account_name,
         api_endpoint=api_endpoint,
@@ -64,11 +65,43 @@ async def _resolve_engine_url(
             raise InterfaceError(f"Unable to retrieve engine endpoint: {e}.")
 
 
+def _validate_engine_name_and_url(
+    engine_name: Optional[str], engine_url: Optional[str]
+) -> None:
+    if engine_name and engine_url:
+        raise ConfigurationError(
+            "Both engine_name and engine_url are provided. Provide only one to connect."
+        )
+    if not engine_name and not engine_url:
+        raise ConfigurationError(
+            "Neither engine_name nor engine_url is provided. Provide one to connect."
+        )
+
+
+def _get_auth(
+    username: Optional[str], password: Optional[str], access_token: Optional[str]
+) -> AuthTypes:
+    if not access_token:
+        if not username or not password:
+            raise ConfigurationError(
+                "Neither username/password nor access_token are provided. Provide one"
+                " to authenticate"
+            )
+        return (username, password)
+    elif username or password:
+        raise ConfigurationError(
+            "Either username/password and access_token are provided. Provide only one"
+            " to authenticate"
+        )
+    return Auth.from_token(access_token)
+
+
 def async_connect_factory(connection_class: Type) -> Callable:
     async def connect_inner(
         database: str = None,
-        username: str = None,
-        password: str = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        access_token: Optional[str] = None,
         engine_name: Optional[str] = None,
         engine_url: Optional[str] = None,
         account_name: Optional[str] = None,
@@ -90,48 +123,31 @@ async def connect_inner(
             Either `engine_name` or `engine_url` should be provided, but not both.
 
         """
-
-        if engine_name and engine_url:
-            raise InterfaceError(
-                "Both engine_name and engine_url are provided. "
-                "Provide only one to connect."
-            )
-        if not engine_name and not engine_url:
-            raise InterfaceError(
-                "Neither engine_name nor engine_url is provided. "
-                "Provide one to connect."
-            )
-
-        api_endpoint = fix_url_schema(api_endpoint)
         # These parameters are optional in function signature
         # but are required to connect.
         # PEP 249 recommends making them kwargs.
-        for param, name in (
-            (database, "database"),
-            (username, "username"),
-            (password, "password"),
-        ):
-            if not param:
-                raise InterfaceError(f"{name} is required to connect.")
+        if not database:
+            raise ConfigurationError("database name is required to connect.")
+
+        _validate_engine_name_and_url(engine_name, engine_url)
+        auth = _get_auth(username, password, access_token)
+        api_endpoint = fix_url_schema(api_endpoint)
 
         # Mypy checks, this should never happen
         assert database is not None
-        assert username is not None
-        assert password is not None
 
         if engine_name:
             engine_url = await _resolve_engine_url(
                 engine_name=engine_name,
-                username=username,
-                password=password,
+                auth=auth,
                 account_name=account_name,
                 api_endpoint=api_endpoint,
             )
 
         assert engine_url is not None
 
         engine_url = fix_url_schema(engine_url)
-        return connection_class(engine_url, database, username, password, api_endpoint)
+        return connection_class(engine_url, database, auth, api_endpoint)
 
     return connect_inner
 
@@ -187,15 +203,16 @@ class BaseConnection:
     def __init__(
         self,
         engine_url: str,
-        database: str,  # TODO: Get by engine name
-        username: str,
-        password: str,
+        database: str,
+        auth: AuthTypes,
         api_endpoint: str = DEFAULT_API_URL,
     ):
+        # Override tcp keepalive settings for connection
         transport = AsyncHTTPTransport()
         transport._pool._network_backend = OverriddenHttpBackend()
+
         self._client = AsyncClient(
-            auth=(username, password),
+            auth=auth,
             base_url=engine_url,
             api_endpoint=api_endpoint,
             timeout=Timeout(DEFAULT_TIMEOUT_SECONDS, read=None),

src/firebolt/client/auth.py

@@ -27,6 +27,12 @@ class Auth(HttpxAuth):
 
     requires_response_body = True
 
+    @staticmethod
+    def from_token(token: str) -> "Auth":
+        a = Auth("", "")
+        a._token = token
+        return a
+
     def __init__(
         self, username: str, password: str, api_endpoint: str = DEFAULT_API_URL
     ):

src/firebolt/common/exception.py

@@ -17,7 +17,7 @@ def __init__(self, method_name: str):
     def __str__(self) -> str:
         return (
             f"Unable to call {self.method_name}: "
-            f"Engine must to be attached to a database first."
+            "Engine must to be attached to a database first."
         )
 
 
@@ -44,7 +44,7 @@ def __init__(self, method_name: str):
     def __str__(self) -> str:
         return (
             f"Unable to call {self.method_name}: "
-            f"Engine must not be in starting or stopping state."
+            "Engine must not be in starting or stopping state."
         )
 
 
@@ -159,3 +159,9 @@ class NotSupportedError(DatabaseError):
     Exception raised when the database encounters an internal error,
     e.g. the cursor is not valid anymore, the transaction is out of sync, etc.
     """
+
+
+class ConfigurationError(InterfaceError):
+    """
+    Exception raised when provided configuration is not correct
+    """

src/firebolt/common/settings.py

@@ -1,12 +1,25 @@
-from pydantic import BaseSettings, Field, SecretStr
+from pydantic import BaseSettings, Field, SecretStr, root_validator
 
 
 class Settings(BaseSettings):
+    # Authorization
+    user: str = Field(None, env="FIREBOLT_USER")
+    password: SecretStr = Field(None, env="FIREBOLT_PASSWORD")
+    # Or
+    access_token: str = Field(None, env="FIREBOLT_AUTH_TOKEN")
+
     account_name: str = Field(None, env="FIREBOLT_ACCOUNT")
     server: str = Field(..., env="FIREBOLT_SERVER")
-    user: str = Field(..., env="FIREBOLT_USER")
-    password: SecretStr = Field(..., env="FIREBOLT_PASSWORD")
     default_region: str = Field(..., env="FIREBOLT_DEFAULT_REGION")
 
     class Config:
         env_file = ".env"
+
+    @root_validator
+    def mutual_exclusive_with_creds(cls, values: dict) -> dict:
+        if values["user"] or values["password"]:
+            if values["access_token"]:
+                raise ValueError("Provide only one of user/password or access_token")
+        elif not values["access_token"]:
+            raise ValueError("Provide either user/password or access_token")
+        return values

src/firebolt/service/manager.py

@@ -1,8 +1,15 @@
 from typing import Optional
 
 from httpx import Timeout
+from httpx._types import AuthTypes
 
-from firebolt.client import Client, log_request, log_response, raise_on_4xx_5xx
+from firebolt.client import (
+    Auth,
+    Client,
+    log_request,
+    log_response,
+    raise_on_4xx_5xx,
+)
 from firebolt.common import Settings
 from firebolt.service.provider import get_provider_id
 
@@ -25,13 +32,17 @@ class ResourceManager:
     - instance types (AWS instance types which engines can use)
     """
 
-    def __init__(
-        self, settings: Optional[Settings] = None, account_name: Optional[str] = None
-    ):
+    def __init__(self, settings: Optional[Settings] = None):
         self.settings = settings or Settings()
 
+        auth: AuthTypes = None
+        if self.settings.access_token:
+            auth = Auth.from_token(self.settings.access_token)
+        else:
+            auth = (self.settings.user, self.settings.password.get_secret_value())
+
         self.client = Client(
-            auth=(self.settings.user, self.settings.password.get_secret_value()),
+            auth=auth,
             base_url=f"https://{ self.settings.server}",
             account_name=self.settings.account_name,
             api_endpoint=self.settings.server,

tests/unit/async_db/test_connection.py

@@ -7,9 +7,9 @@
 from firebolt.async_db import Connection, connect
 from firebolt.async_db._types import ColType
 from firebolt.common.exception import (
+    ConfigurationError,
     ConnectionClosedError,
     FireboltEngineError,
-    InterfaceError,
 )
 from firebolt.common.settings import Settings
 from firebolt.common.urls import ACCOUNT_ENGINE_BY_NAME_URL
@@ -90,18 +90,49 @@ async def test_cursor_initialized(
 
 @mark.asyncio
 async def test_connect_empty_parameters():
-    params = ("database", "username", "password")
-    kwargs = {"engine_url": "engine_url", **{p: p for p in params}}
+    with raises(ConfigurationError):
+        async with await connect(engine_url="engine_url"):
+            pass
+
+
+@mark.asyncio
+async def test_connect_access_token(
+    settings: Settings,
+    db_name: str,
+    httpx_mock: HTTPXMock,
+    auth_callback: Callable,
+    auth_url: str,
+    check_token_callback: Callable,
+    query_url: str,
+    python_query_data: List[List[ColType]],
+    access_token: str,
+):
+    httpx_mock.add_callback(check_token_callback, url=query_url)
+    async with (
+        await connect(
+            engine_url=settings.server,
+            database=db_name,
+            access_token=access_token,
+            account_name="a",
+            api_endpoint=settings.server,
+        )
+    ) as connection:
+        cursor = connection.cursor()
+        assert await cursor.execute("select*") == -1
+
+    with raises(ConfigurationError):
+        async with await connect(engine_url="engine_url", database="database"):
+            pass
 
-    for param in params:
-        with raises(InterfaceError) as exc_info:
-            kwargs = {
-                "engine_url": "engine_url",
-                **{p: p for p in params if p != param},
-            }
-            async with await connect(**kwargs):
-                pass
-        assert str(exc_info.value) == f"{param} is required to connect."
+    with raises(ConfigurationError):
+        async with await connect(
+            engine_url="engine_url",
+            database="database",
+            username="username",
+            password="password",
+            access_token="access_token",
+        ):
+            pass
 
 
 @mark.asyncio
@@ -125,7 +156,7 @@ async def test_connect_engine_name(
 ):
     """connect properly handles engine_name"""
 
-    with raises(InterfaceError) as exc_info:
+    with raises(ConfigurationError):
         async with await connect(
             engine_url="engine_url",
             engine_name="engine_name",
@@ -135,21 +166,15 @@ async def test_connect_engine_name(
             account_name="account_name",
         ):
             pass
-    assert str(exc_info.value).startswith(
-        "Both engine_name and engine_url are provided."
-    )
 
-    with raises(InterfaceError) as exc_info:
+    with raises(ConfigurationError):
         async with await connect(
             database="db",
             username="username",
             password="password",
             account_name="account",
         ):
             pass
-    assert str(exc_info.value).startswith(
-        "Neither engine_name nor engine_url is provided."
-    )
 
     httpx_mock.add_callback(auth_callback, url=auth_url)
     httpx_mock.add_callback(query_callback, url=query_url)
@@ -166,7 +191,7 @@ async def test_connect_engine_name(
         status_code=codes.NOT_FOUND,
     )
 
-    with raises(FireboltEngineError) as exc_info:
+    with raises(FireboltEngineError):
         async with await connect(
             database="db",
             username="username",