Skip to content

Commit e702dc8

Browse files
ptiurinptiurin
authored
feat: Service account authentication (#210)
1 parent c71d252 commit e702dc8

File tree

15 files changed

+318
-74
lines changed

15 files changed

+318
-74
lines changed

‎.github/workflows/integration-tests.yml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@ on:
77
required: true
88
FIREBOLT_PASSWORD:
99
required: true
10+
SERVICE_ID:
11+
required: true
12+
SERVICE_SECRET:
13+
required: true
1014
jobs:
1115
tests:
1216
runs-on: ubuntu-latest
@@ -37,6 +41,8 @@ jobs:
3741
env:
3842
USER_NAME: ${{ secrets.FIREBOLT_USERNAME }}
3943
PASSWORD: ${{ secrets.FIREBOLT_PASSWORD }}
44+
SERVICE_ID: ${{ secrets.SERVICE_ID }}
45+
SERVICE_SECRET: ${{ secrets.SERVICE_SECRET }}
4046
DATABASE_NAME: ${{ steps.setup.outputs.database_name }}
4147
ENGINE_NAME: ${{ steps.setup.outputs.engine_name }}
4248
ENGINE_URL: ${{ steps.setup.outputs.engine_url }}

‎.github/workflows/nightly.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,8 @@ jobs:
5656
env:
5757
USER_NAME: ${{ secrets.FIREBOLT_USERNAME }}
5858
PASSWORD: ${{ secrets.FIREBOLT_PASSWORD }}
59+
SERVICE_ID: ${{ secrets.SERVICE_ID }}
60+
SERVICE_SECRET: ${{ secrets.SERVICE_SECRET }}
5961
DATABASE_NAME: ${{ steps.setup.outputs.database_name }}
6062
ENGINE_NAME: ${{ steps.setup.outputs.engine_name }}
6163
ENGINE_URL: ${{ steps.setup.outputs.engine_url }}

‎.github/workflows/release.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ jobs:
1616
secrets:
1717
FIREBOLT_USERNAME: ${{ secrets.FIREBOLT_USERNAME }}
1818
FIREBOLT_PASSWORD: ${{ secrets.FIREBOLT_PASSWORD }}
19+
SERVICE_ID: ${{ secrets.SERVICE_ID }}
20+
SERVICE_SECRET: ${{ secrets.SERVICE_SECRET }}
1921

2022
publish:
2123
runs-on: ubuntu-latest

‎src/firebolt/client/auth/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
from firebolt.client.auth.base import Auth
2+
from firebolt.client.auth.service_account import ServiceAccount
23
from firebolt.client.auth.token import Token
34
from firebolt.client.auth.username_password import UsernamePassword

‎src/firebolt/client/auth/request_auth_base.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
from time import time
2+
from typing import Generator
3+
4+
from httpx import Request, Response
5+
6+
from firebolt.client.auth.base import Auth
7+
from firebolt.client.constants import _REQUEST_ERRORS
8+
from firebolt.utils.exception import AuthenticationError
9+
from firebolt.utils.usage_tracker import get_user_agent_header
10+
11+
12+
class _RequestBasedAuth(Auth):
13+
"""Base abstract class for http request based authentication."""
14+
15+
def __init__(self, use_token_cache: bool = True):
16+
self._user_agent = get_user_agent_header()
17+
super().__init__(use_token_cache)
18+
19+
def _make_auth_request(self) -> Request:
20+
"""Create an HTTP request required for authentication.
21+
Returns:
22+
Request: HTTP request, required for authentication.
23+
"""
24+
raise NotImplementedError()
25+
26+
@staticmethod
27+
def _check_response_error(response: dict) -> None:
28+
"""Check if response data contains errors.
29+
Args:
30+
response (dict): Response data
31+
Raises:
32+
AuthenticationError: Were unable to authenticate
33+
"""
34+
if "error" in response:
35+
raise AuthenticationError(
36+
response.get("message", "unknown server error"),
37+
)
38+
39+
def get_new_token_generator(self) -> Generator[Request, Response, None]:
40+
"""Get new token using username and password.
41+
Yields:
42+
Request: An http request to get token. Expects Response to be sent back
43+
Raises:
44+
AuthenticationError: Error while authenticating with provided credentials
45+
"""
46+
try:
47+
response = yield self._make_auth_request()
48+
response.raise_for_status()
49+
50+
parsed = response.json()
51+
self._check_response_error(parsed)
52+
53+
self._token = parsed["access_token"]
54+
self._expires = int(time()) + int(parsed["expires_in"])
55+
56+
except _REQUEST_ERRORS as e:
57+
raise AuthenticationError(repr(e)) from e

‎src/firebolt/client/auth/service_account.py‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
from typing import Optional
2+
3+
from firebolt.client.auth.base import AuthRequest
4+
from firebolt.client.auth.request_auth_base import _RequestBasedAuth
5+
from firebolt.utils.token_storage import TokenSecureStorage
6+
from firebolt.utils.urls import AUTH_SERVICE_ACCOUNT_URL
7+
from firebolt.utils.util import cached_property
8+
9+
10+
class ServiceAccount(_RequestBasedAuth):
11+
"""Service Account authentication class for Firebolt Database.
12+
13+
Gets authentication token using
14+
provided credentials and updates it when it expires.
15+
16+
Args:
17+
client_id (str): Client ID
18+
client_secret (str): Client secret
19+
use_token_cache (bool): True if token should be cached in filesystem;
20+
False otherwise
21+
22+
Attributes:
23+
client_id (str): Client ID
24+
client_secret (str): Client secret
25+
"""
26+
27+
__slots__ = (
28+
"client_id",
29+
"client_secret",
30+
"_token",
31+
"_expires",
32+
"_use_token_cache",
33+
"_user_agent",
34+
)
35+
36+
requires_response_body = True
37+
38+
def __init__(
39+
self,
40+
client_id: str,
41+
client_secret: str,
42+
use_token_cache: bool = True,
43+
):
44+
self.client_id = client_id
45+
self.client_secret = client_secret
46+
super().__init__(use_token_cache)
47+
48+
def copy(self) -> "ServiceAccount":
49+
"""Make another auth object with same credentials.
50+
51+
Returns:
52+
ServiceAccount: Auth object
53+
"""
54+
return ServiceAccount(self.client_id, self.client_secret, self._use_token_cache)
55+
56+
@cached_property
57+
def _token_storage(self) -> Optional[TokenSecureStorage]:
58+
"""Token filesystem cache storage.
59+
60+
This is evaluated lazily, only if caching is enabled
61+
62+
Returns:
63+
TokenSecureStorage: Token filesystem cache storage
64+
"""
65+
return TokenSecureStorage(username=self.client_id, password=self.client_secret)
66+
67+
def _make_auth_request(self) -> AuthRequest:
68+
"""Get new token using username and password.
69+
70+
Yields:
71+
Request: An http request to get token. Expects Response to be sent back
72+
73+
Raises:
74+
AuthenticationError: Error while authenticating with provided credentials
75+
"""
76+
77+
response = self.request_class(
78+
"POST",
79+
AUTH_SERVICE_ACCOUNT_URL,
80+
headers={
81+
"User-Agent": self._user_agent,
82+
},
83+
data={
84+
"client_id": self.client_id,
85+
"client_secret": self.client_secret,
86+
"grant_type": "client_credentials",
87+
},
88+
)
89+
return response

‎src/firebolt/client/auth/username_password.py‎

Lines changed: 16 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,13 @@
1-
from time import time
2-
from typing import Generator, Optional
1+
from typing import Optional
32

4-
from httpx import Request, Response
5-
6-
from firebolt.client.auth.base import Auth
7-
from firebolt.client.constants import _REQUEST_ERRORS
8-
from firebolt.utils.exception import AuthenticationError
3+
from firebolt.client.auth.base import AuthRequest
4+
from firebolt.client.auth.request_auth_base import _RequestBasedAuth
95
from firebolt.utils.token_storage import TokenSecureStorage
106
from firebolt.utils.urls import AUTH_URL
117
from firebolt.utils.util import cached_property
128

139

14-
class UsernamePassword(Auth):
10+
class UsernamePassword(_RequestBasedAuth):
1511
"""Username/Password authentication class for Firebolt Database.
1612
1713
Gets authentication token using
@@ -34,6 +30,7 @@ class UsernamePassword(Auth):
3430
"_token",
3531
"_expires",
3632
"_use_token_cache",
33+
"_user_agent",
3734
)
3835

3936
requires_response_body = True
@@ -67,7 +64,7 @@ def _token_storage(self) -> Optional[TokenSecureStorage]:
6764
"""
6865
return TokenSecureStorage(username=self.username, password=self.password)
6966

70-
def get_new_token_generator(self) -> Generator[Request, Response, None]:
67+
def _make_auth_request(self) -> AuthRequest:
7168
"""Get new token using username and password.
7269
7370
Yields:
@@ -76,39 +73,13 @@ def get_new_token_generator(self) -> Generator[Request, Response, None]:
7673
Raises:
7774
AuthenticationError: Error while authenticating with provided credentials
7875
"""
79-
try:
80-
response = yield self.request_class(
81-
"POST",
82-
# The full url is generated on client side by attaching
83-
# it to api_endpoint
84-
AUTH_URL,
85-
headers={
86-
"Content-Type": "application/json;charset=UTF-8",
87-
"User-Agent": "firebolt-sdk",
88-
},
89-
json={"username": self.username, "password": self.password},
90-
)
91-
response.raise_for_status()
92-
93-
parsed = response.json()
94-
self._check_response_error(parsed)
95-
96-
self._token = parsed["access_token"]
97-
self._expires = int(time()) + int(parsed["expires_in"])
98-
99-
except _REQUEST_ERRORS as e:
100-
raise AuthenticationError(repr(e)) from e
101-
102-
def _check_response_error(self, response: dict) -> None:
103-
"""Check if response data contains errors.
104-
105-
Args:
106-
response (dict): Response data
107-
108-
Raises:
109-
AuthenticationError: Were unable to authenticate
110-
"""
111-
if "error" in response:
112-
raise AuthenticationError(
113-
response.get("message", "unknown server error"),
114-
)
76+
response = self.request_class(
77+
"POST",
78+
AUTH_URL,
79+
headers={
80+
"Content-Type": "application/json;charset=UTF-8",
81+
"User-Agent": self._user_agent,
82+
},
83+
json={"username": self.username, "password": self.password},
84+
)
85+
return response

‎src/firebolt/utils/urls.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
AUTH_URL = "/auth/v1/login"
2+
AUTH_SERVICE_ACCOUNT_URL = "/auth/v1/token"
23

34
DATABASES_URL = "/core/v1/account/databases"
45

‎tests/integration/conftest.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
PASSWORD_ENV = "PASSWORD"
1818
ACCOUNT_NAME_ENV = "ACCOUNT_NAME"
1919
API_ENDPOINT_ENV = "API_ENDPOINT"
20+
SERVICE_ID_ENV = "SERVICE_ID"
21+
SERVICE_SECRET_ENV = "SERVICE_SECRET"
2022

2123

2224
def must_env(var_name: str) -> str:
@@ -78,3 +80,13 @@ def account_name() -> str:
7880
@fixture(scope="session")
7981
def api_endpoint() -> str:
8082
return must_env(API_ENDPOINT_ENV)
83+
84+
85+
@fixture(scope="session")
86+
def service_id() -> str:
87+
return must_env(SERVICE_ID_ENV)
88+
89+
90+
@fixture(scope="session")
91+
def service_secret() -> str:
92+
return must_env(SERVICE_SECRET_ENV)

‎tests/integration/dbapi/async/conftest.py‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from pytest_asyncio import fixture as async_fixture
22

33
from firebolt.async_db import Connection, connect
4+
from firebolt.client.auth import ServiceAccount, UsernamePassword
45

56

67
@async_fixture
@@ -15,8 +16,26 @@ async def connection(
1516
async with await connect(
1617
engine_url=engine_url,
1718
database=database_name,
18-
username=username,
19-
password=password,
19+
auth=UsernamePassword(username, password),
20+
account_name=account_name,
21+
api_endpoint=api_endpoint,
22+
) as connection:
23+
yield connection
24+
25+
26+
@async_fixture
27+
async def service_account_connection(
28+
engine_url: str,
29+
database_name: str,
30+
service_id: str,
31+
service_secret: str,
32+
account_name: str,
33+
api_endpoint: str,
34+
) -> Connection:
35+
async with await connect(
36+
engine_url=engine_url,
37+
database=database_name,
38+
auth=ServiceAccount(service_id, service_secret),
2039
account_name=account_name,
2140
api_endpoint=api_endpoint,
2241
) as connection:

0 commit comments

Comments
 (0)