From 5aa075addee03f1ede77b16a64255abed77ca8ce Mon Sep 17 00:00:00 2001
From: ptiurin <petro.tiurin@firebolt.io>
Date: Mon, 4 Dec 2023 14:01:55 +0000
Subject: [PATCH] better auth detection

---
 .github/workflows/integration-tests-v2.yml | 16 ++++++++--------
 dbt/adapters/firebolt/connections.py       |  8 ++++++--
 tests/conftest.py                          | 14 +++++++++++---
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/integration-tests-v2.yml b/.github/workflows/integration-tests-v2.yml
index 60c14b841..626c1cd59 100644
--- a/.github/workflows/integration-tests-v2.yml
+++ b/.github/workflows/integration-tests-v2.yml
@@ -47,11 +47,11 @@ jobs:
       - name: Determine env variables
         run: |
           if [ "${{ inputs.environment }}" == 'staging' ]; then
-             echo "KEY=${{ secrets.FIREBOLT_CLIENT_ID_STG_NEW_IDN }}" >> "$GITHUB_ENV"
-             echo "SECRET=${{ secrets.FIREBOLT_CLIENT_SECRET_STG_NEW_IDN }}" >> "$GITHUB_ENV"
+             echo "CLIENT_ID=${{ secrets.FIREBOLT_CLIENT_ID_STG_NEW_IDN }}" >> "$GITHUB_ENV"
+             echo "CLIENT_SECRET=${{ secrets.FIREBOLT_CLIENT_SECRET_STG_NEW_IDN }}" >> "$GITHUB_ENV"
           else
-             echo "KEY=${{ secrets.FIREBOLT_CLIENT_ID_DEV_NEW_IDN }}" >> "$GITHUB_ENV"
-             echo "SECRET=${{ secrets.FIREBOLT_CLIENT_SECRET_DEV_NEW_IDN }}" >> "$GITHUB_ENV"
+             echo "CLIENT_ID=${{ secrets.FIREBOLT_CLIENT_ID_DEV_NEW_IDN }}" >> "$GITHUB_ENV"
+             echo "CLIENT_SECRET=${{ secrets.FIREBOLT_CLIENT_SECRET_DEV_NEW_IDN }}" >> "$GITHUB_ENV"
           fi
 
       - name: Keep environment name in the summary
@@ -61,8 +61,8 @@ jobs:
         id: setup
         uses: firebolt-db/integration-testing-setup@v2
         with:
-          firebolt-client-id: ${{ env.KEY }}
-          firebolt-client-secret: ${{ env.SECRET }}
+          firebolt-client-id: ${{ env.CLIENT_ID }}
+          firebolt-client-secret: ${{ env.CLIENT_SECRET }}
           account: "automation"
           api-endpoint: "api.${{ inputs.environment }}.firebolt.io"
 
@@ -76,8 +76,8 @@ jobs:
 
       - name: Run integration tests
         env:
-          USER_NAME: ${{ env.KEY }}
-          PASSWORD: ${{ env.SECRET }}
+          CLIENT_ID: ${{ env.CLIENT_ID }}
+          CLIENT_SECRET: ${{ env.CLIENT_SECRET }}
           DATABASE_NAME: ${{ steps.setup.outputs.database_name }}
           ENGINE_NAME: ${{ steps.setup.outputs.engine_name }}
           API_ENDPOINT: "api.${{ inputs.environment }}.firebolt.io"
diff --git a/dbt/adapters/firebolt/connections.py b/dbt/adapters/firebolt/connections.py
index bfdd4193b..01b3f610d 100644
--- a/dbt/adapters/firebolt/connections.py
+++ b/dbt/adapters/firebolt/connections.py
@@ -95,10 +95,14 @@ def open(cls, connection: Connection) -> Connection:
             return connection
         credentials = connection.credentials
         auth: Auth  # mypy typecheck
-        # check if we're using email and use the appropriate auth
-        if '@' in credentials.user:
+        # client id and client secret are the new way to authenticate
+        if hasattr(credentials, 'client_id') and hasattr(credentials, 'client_secret'):
+            auth = ClientCredentials(credentials.client_id, credentials.client_secret)
+        elif '@' in credentials.user:
+            # email auth can only be used with UsernamePassword
             auth = UsernamePassword(credentials.user, credentials.password)
         else:
+            # assume user provided id and secret in the user/password fields
             auth = ClientCredentials(credentials.user, credentials.password)
 
         def connect() -> SDKConnection:
diff --git a/tests/conftest.py b/tests/conftest.py
index d765960c2..18ec887d2 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -11,17 +11,25 @@
 # dbt will supply a unique schema per test, so we do not specify 'schema' here
 @pytest.fixture(scope='class')
 def dbt_profile_target():
-    return {
+    profile = {
         'type': 'firebolt',
         'threads': 2,
         'api_endpoint': os.getenv('API_ENDPOINT'),
         'account_name': os.getenv('ACCOUNT_NAME'),
         'database': os.getenv('DATABASE_NAME'),
         'engine_name': os.getenv('ENGINE_NAME'),
-        'user': os.getenv('USER_NAME'),
-        'password': os.getenv('PASSWORD'),
         'port': 443,
     }
+    # add credentials to the profile keys
+    if os.getenv('USER_NAME') and os.getenv('PASSWORD'):
+        profile['user'] = os.getenv('USER_NAME')
+        profile['password'] = os.getenv('PASSWORD')
+    elif os.getenv('CLIENT_ID') and os.getenv('CLIENT_SECRET'):
+        profile['client_id'] = os.getenv('CLIENT_ID')
+        profile['client_secret'] = os.getenv('CLIENT_SECRET')
+    else:
+        raise Exception('No credentials found in environment')
+    return profile
 
 
 # Overriding dbt_profile_data in order to set the schema to public.