From fc7301da513971e7a231e609637ba0d2770b0125 Mon Sep 17 00:00:00 2001 From: ptiurin Date: Wed, 7 Dec 2022 15:24:13 +0000 Subject: [PATCH 1/7] feat: Service account support --- src/firebolt_db/firebolt_dialect.py | 10 +++++-- tests/integration/conftest.py | 28 +++++++++++++++++++ .../test_sqlalchemy_integration.py | 4 +++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/src/firebolt_db/firebolt_dialect.py b/src/firebolt_db/firebolt_dialect.py index c94270a..2d05801 100644 --- a/src/firebolt_db/firebolt_dialect.py +++ b/src/firebolt_db/firebolt_dialect.py @@ -5,7 +5,7 @@ import firebolt.db as dbapi import sqlalchemy.types as sqltypes -from firebolt.client.auth import Auth, UsernamePassword +from firebolt.client.auth import Auth, ServiceAccount, UsernamePassword from firebolt.db import Cursor from sqlalchemy.engine import Connection as AlchemyConnection from sqlalchemy.engine import ExecutionContext, default @@ -111,9 +111,15 @@ def create_connect_args(self, url: URL) -> Tuple[List, Dict]: # parameters are all passed as a string, we need to convert # bool flag to boolean for SDK compatibility token_cache_flag = bool(strtobool(parameters.pop("use_token_cache", "True"))) + service_account_flag = bool(strtobool(parameters.pop("service_account", "0"))) + auth = ( + ServiceAccount(url.username, url.password, token_cache_flag) + if service_account_flag + else UsernamePassword(url.username, url.password, token_cache_flag) + ) kwargs: Dict[str, Union[str, Auth, Dict[str, Any], None]] = { "database": url.host or None, - "auth": UsernamePassword(url.username, url.password, token_cache_flag), + "auth": auth, "engine_name": url.database, "additional_parameters": {}, } diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index b321b36..f0ef581 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -13,6 +13,8 @@ DATABASE_NAME_ENV = "DATABASE_NAME" USERNAME_ENV = "USER_NAME" PASSWORD_ENV = "PASSWORD" +SERVICE_ID = "SERVICE_ID" +SERVICE_SECRET = "SERVICE_SECRET" def must_env(var_name: str) -> str: @@ -41,6 +43,16 @@ def password() -> str: return must_env(PASSWORD_ENV) +@fixture(scope="session") +def service_id() -> str: + return must_env(SERVICE_ID) + + +@fixture(scope="session") +def service_secret() -> str: + return must_env(SERVICE_SECRET) + + @fixture(scope="session") def engine( username: str, password: str, database_name: str, engine_name: str @@ -50,12 +62,28 @@ def engine( ) +@fixture(scope="session") +def engine_service_account( + service_id: str, service_secret: str, database_name: str, engine_name: str +) -> Engine: + return create_engine( + f"firebolt://{service_id}:{service_secret}@{database_name}/{engine_name}" + "?service_account=1" + ) + + @fixture(scope="session") def connection(engine: Engine) -> Connection: with engine.connect() as c: yield c +@fixture(scope="session") +def connection_service_account(engine_service_account: Engine) -> Connection: + with engine_service_account.connect() as c: + yield c + + @fixture(scope="session") def event_loop(): loop = asyncio.get_event_loop() diff --git a/tests/integration/test_sqlalchemy_integration.py b/tests/integration/test_sqlalchemy_integration.py index 5af0198..a43f67d 100644 --- a/tests/integration/test_sqlalchemy_integration.py +++ b/tests/integration/test_sqlalchemy_integration.py @@ -111,3 +111,7 @@ def test_get_columns(self, engine: Engine, fact_table_name: str): assert row_keys[1] == "type" assert row_keys[2] == "nullable" assert row_keys[3] == "default" + + def test_service_account_connect(self, connection_service_account: Connection): + result = connection_service_account.execute("SELECT 1") + assert result.fetchall() == [(1,)] From 0ff95da80cc9e870dd7296b458bbe22e5b6eda87 Mon Sep 17 00:00:00 2001 From: ptiurin Date: Wed, 7 Dec 2022 16:12:03 +0000 Subject: [PATCH 2/7] Fixing integration test secrets --- .github/workflows/nightly.yml | 2 ++ .github/workflows/python-integration-tests.yml | 2 ++ .github/workflows/release.yml | 2 ++ 3 files changed, 6 insertions(+) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 982fdc5..8268684 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -56,6 +56,8 @@ jobs: env: USER_NAME: ${{ secrets.FIREBOLT_USERNAME }} PASSWORD: ${{ secrets.FIREBOLT_PASSWORD }} + SERVICE_ID: ${{ secrets.SERVICE_ID }} + SERVICE_SECRET: ${{ secrets.SERVICE_SECRET }} DATABASE_NAME: ${{ steps.setup.outputs.database_name }} ENGINE_NAME: ${{ steps.setup.outputs.engine_name }} ENGINE_URL: ${{ steps.setup.outputs.engine_url }} diff --git a/.github/workflows/python-integration-tests.yml b/.github/workflows/python-integration-tests.yml index 3080b31..aa40b8c 100644 --- a/.github/workflows/python-integration-tests.yml +++ b/.github/workflows/python-integration-tests.yml @@ -33,6 +33,8 @@ jobs: env: USER_NAME: ${{ secrets.FIREBOLT_USERNAME }} PASSWORD: ${{ secrets.FIREBOLT_PASSWORD }} + SERVICE_ID: ${{ secrets.SERVICE_ID }} + SERVICE_SECRET: ${{ secrets.SERVICE_SECRET }} DATABASE_NAME: ${{ steps.setup.outputs.database_name }} ENGINE_NAME: ${{ steps.setup.outputs.engine_name }} ENGINE_URL: ${{ steps.setup.outputs.engine_url }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 91d09c5..cf03584 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -40,6 +40,8 @@ jobs: env: USER_NAME: ${{ secrets.FIREBOLT_USERNAME }} PASSWORD: ${{ secrets.FIREBOLT_PASSWORD }} + SERVICE_ID: ${{ secrets.SERVICE_ID }} + SERVICE_SECRET: ${{ secrets.SERVICE_SECRET }} DATABASE_NAME: ${{ steps.setup.outputs.database_name }} ENGINE_NAME: ${{ steps.setup.outputs.engine_name }} ENGINE_URL: ${{ steps.setup.outputs.engine_url }} From 8eb77b96b58d272ba1097e15fc9122b57684a47c Mon Sep 17 00:00:00 2001 From: ptiurin Date: Fri, 9 Dec 2022 14:35:55 +0000 Subject: [PATCH 3/7] Auto-detecting auth --- src/firebolt_db/firebolt_dialect.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/firebolt_db/firebolt_dialect.py b/src/firebolt_db/firebolt_dialect.py index 2d05801..10b4832 100644 --- a/src/firebolt_db/firebolt_dialect.py +++ b/src/firebolt_db/firebolt_dialect.py @@ -111,10 +111,9 @@ def create_connect_args(self, url: URL) -> Tuple[List, Dict]: # parameters are all passed as a string, we need to convert # bool flag to boolean for SDK compatibility token_cache_flag = bool(strtobool(parameters.pop("use_token_cache", "True"))) - service_account_flag = bool(strtobool(parameters.pop("service_account", "0"))) auth = ( ServiceAccount(url.username, url.password, token_cache_flag) - if service_account_flag + if "@" in url.username else UsernamePassword(url.username, url.password, token_cache_flag) ) kwargs: Dict[str, Union[str, Auth, Dict[str, Any], None]] = { From 52928775eca4a8c28370fc12dceb7c0f9cd09c34 Mon Sep 17 00:00:00 2001 From: ptiurin Date: Mon, 12 Dec 2022 10:01:58 +0000 Subject: [PATCH 4/7] back to query parameter --- src/firebolt_db/firebolt_dialect.py | 5 +++- tests/unit/test_firebolt_dialect.py | 37 ++++++++++++++++++++++++++--- 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/src/firebolt_db/firebolt_dialect.py b/src/firebolt_db/firebolt_dialect.py index 10b4832..c96630b 100644 --- a/src/firebolt_db/firebolt_dialect.py +++ b/src/firebolt_db/firebolt_dialect.py @@ -111,9 +111,12 @@ def create_connect_args(self, url: URL) -> Tuple[List, Dict]: # parameters are all passed as a string, we need to convert # bool flag to boolean for SDK compatibility token_cache_flag = bool(strtobool(parameters.pop("use_token_cache", "True"))) + service_account_flag = bool( + strtobool(parameters.pop("service_account", "False")) + ) auth = ( ServiceAccount(url.username, url.password, token_cache_flag) - if "@" in url.username + if service_account_flag else UsernamePassword(url.username, url.password, token_cache_flag) ) kwargs: Dict[str, Union[str, Auth, Dict[str, Any], None]] = { diff --git a/tests/unit/test_firebolt_dialect.py b/tests/unit/test_firebolt_dialect.py index a53c71d..377ac13 100644 --- a/tests/unit/test_firebolt_dialect.py +++ b/tests/unit/test_firebolt_dialect.py @@ -30,11 +30,42 @@ def test_create_dialect(self, dialect: FireboltDialect): assert isinstance(dialect.type_compiler, FireboltTypeCompiler) assert dialect.context == {} - def test_create_connect_args(self, dialect: FireboltDialect): + @mark.parametrize( + "query_params", + [ + ("service_account=True"), + ("service_account=true"), + ("service_account=1"), + ], + ) + def test_create_connect_args_service_account( + self, dialect: FireboltDialect, query_params: str + ): + u = url.make_url( + "test_engine://test-sa-user-key:test_password@test_db_name/test_engine_name" + + "?" + + query_params + ) + with mock.patch.dict(os.environ, {"FIREBOLT_BASE_URL": "test_url"}): + result_list, result_dict = dialect.create_connect_args(u) + assert result_dict["engine_name"] == "test_engine_name" + assert result_dict["auth"].client_id == "test-sa-user-key" + assert result_dict["auth"].client_secret == "test_password" + assert result_dict["auth"]._use_token_cache is True + assert result_dict["database"] == "test_db_name" + assert result_dict["api_endpoint"] == "test_url" + assert "username" not in result_dict + assert "password" not in result_dict + assert result_list == [] + + @mark.parametrize( + "query_params", [(""), ("service_account=0"), ("service_account=False")] + ) + def test_create_connect_args(self, dialect: FireboltDialect, query_params: str): connection_url = ( - "test_engine://test_user@email:test_password@test_db_name/test_engine_name" + "test_engine://test_user@email:test_password@test_db_name/test_engine_name?" ) - u = url.make_url(connection_url) + u = url.make_url(connection_url + query_params) with mock.patch.dict(os.environ, {"FIREBOLT_BASE_URL": "test_url"}): result_list, result_dict = dialect.create_connect_args(u) assert result_dict["engine_name"] == "test_engine_name" From 443165d30bdcc12b2736301c543d4a8cae0b0bbb Mon Sep 17 00:00:00 2001 From: ptiurin Date: Mon, 12 Dec 2022 10:22:01 +0000 Subject: [PATCH 5/7] legitignore --- .legitignore | 1 + 1 file changed, 1 insertion(+) create mode 100644 .legitignore diff --git a/.legitignore b/.legitignore new file mode 100644 index 0000000..8b109ac --- /dev/null +++ b/.legitignore @@ -0,0 +1 @@ +tests/unit/** # Ignore everything in the unit test directory From ac5f72d2b0b2ba349b1ac8e068a29716d20b9d4e Mon Sep 17 00:00:00 2001 From: ptiurin Date: Mon, 12 Dec 2022 11:55:02 +0000 Subject: [PATCH 6/7] @ is working now... --- src/firebolt_db/firebolt_dialect.py | 5 +---- tests/unit/test_firebolt_dialect.py | 21 +++------------------ 2 files changed, 4 insertions(+), 22 deletions(-) diff --git a/src/firebolt_db/firebolt_dialect.py b/src/firebolt_db/firebolt_dialect.py index c96630b..192259f 100644 --- a/src/firebolt_db/firebolt_dialect.py +++ b/src/firebolt_db/firebolt_dialect.py @@ -111,12 +111,9 @@ def create_connect_args(self, url: URL) -> Tuple[List, Dict]: # parameters are all passed as a string, we need to convert # bool flag to boolean for SDK compatibility token_cache_flag = bool(strtobool(parameters.pop("use_token_cache", "True"))) - service_account_flag = bool( - strtobool(parameters.pop("service_account", "False")) - ) auth = ( ServiceAccount(url.username, url.password, token_cache_flag) - if service_account_flag + if "@" not in url.username else UsernamePassword(url.username, url.password, token_cache_flag) ) kwargs: Dict[str, Union[str, Auth, Dict[str, Any], None]] = { diff --git a/tests/unit/test_firebolt_dialect.py b/tests/unit/test_firebolt_dialect.py index 377ac13..e08f89c 100644 --- a/tests/unit/test_firebolt_dialect.py +++ b/tests/unit/test_firebolt_dialect.py @@ -30,21 +30,9 @@ def test_create_dialect(self, dialect: FireboltDialect): assert isinstance(dialect.type_compiler, FireboltTypeCompiler) assert dialect.context == {} - @mark.parametrize( - "query_params", - [ - ("service_account=True"), - ("service_account=true"), - ("service_account=1"), - ], - ) - def test_create_connect_args_service_account( - self, dialect: FireboltDialect, query_params: str - ): + def test_create_connect_args_service_account(self, dialect: FireboltDialect): u = url.make_url( "test_engine://test-sa-user-key:test_password@test_db_name/test_engine_name" - + "?" - + query_params ) with mock.patch.dict(os.environ, {"FIREBOLT_BASE_URL": "test_url"}): result_list, result_dict = dialect.create_connect_args(u) @@ -58,14 +46,11 @@ def test_create_connect_args_service_account( assert "password" not in result_dict assert result_list == [] - @mark.parametrize( - "query_params", [(""), ("service_account=0"), ("service_account=False")] - ) - def test_create_connect_args(self, dialect: FireboltDialect, query_params: str): + def test_create_connect_args(self, dialect: FireboltDialect): connection_url = ( "test_engine://test_user@email:test_password@test_db_name/test_engine_name?" ) - u = url.make_url(connection_url + query_params) + u = url.make_url(connection_url) with mock.patch.dict(os.environ, {"FIREBOLT_BASE_URL": "test_url"}): result_list, result_dict = dialect.create_connect_args(u) assert result_dict["engine_name"] == "test_engine_name" From 0d5be4501ecb4bdd2058c1939d94de5c340c00e2 Mon Sep 17 00:00:00 2001 From: ptiurin Date: Mon, 12 Dec 2022 14:39:54 +0000 Subject: [PATCH 7/7] Fix conftest --- tests/integration/conftest.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index f0ef581..ac1c031 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -68,7 +68,6 @@ def engine_service_account( ) -> Engine: return create_engine( f"firebolt://{service_id}:{service_secret}@{database_name}/{engine_name}" - "?service_account=1" )