From 46f0559efb58d599e436e7218163758218c129a0 Mon Sep 17 00:00:00 2001 From: Pradeep Srikakolapu Date: Mon, 25 Nov 2024 18:41:43 -0800 Subject: [PATCH] Testing OIDC --- .github/workflows/integration-tests-azure.yml | 103 ++++++++++-------- .../fabric/fabric_connection_manager.py | 40 ++++++- dbt/adapters/fabric/fabric_credentials.py | 1 + tests/conftest.py | 19 +++- 4 files changed, 114 insertions(+), 49 deletions(-) diff --git a/.github/workflows/integration-tests-azure.yml b/.github/workflows/integration-tests-azure.yml index 4b3da67..930cd3a 100644 --- a/.github/workflows/integration-tests-azure.yml +++ b/.github/workflows/integration-tests-azure.yml @@ -1,13 +1,10 @@ ---- name: Integration tests on Fabric DW -on: # yamllint disable-line rule:truthy +on: # yamllint disable-line rule:truthy workflow_dispatch: pull_request: branches: - oidc_connect - - jobs: integration-tests-fabric-dw: name: Regular @@ -15,9 +12,9 @@ jobs: fail-fast: false max-parallel: 1 matrix: - profile: ["ci_azure_auto"] + profile: ["integration_tests"] python_version: ["3.11"] - msodbc_version: ["17", "18"] + msodbc_version: ["18"] runs-on: ubuntu-latest permissions: @@ -27,49 +24,65 @@ jobs: container: image: ghcr.io/${{ github.repository }}:CI-${{ matrix.python_version }}-msodbc${{ matrix.msodbc_version }} steps: - # Azure login using federated credentials - - name: Azure login with OIDC - uses: azure/login@v2 - with: - client-id: ${{ secrets.DBT_AZURE_SP_NAME }} - tenant-id: ${{ secrets.DBT_AZURE_TENANT }} - allow-no-subscriptions: true - federated-token: true + # Azure login using federated credentials + - name: Azure login with OIDC + uses: azure/login@v2 + with: + client-id: ${{ secrets.DBT_AZURE_SP_NAME }} + tenant-id: ${{ secrets.DBT_AZURE_TENANT }} + allow-no-subscriptions: true + federated-token: true + + - name: Connect to Fabric Warehouse to Retrieve Token + id: fetch_token + run: | + pip install azure-identity pyodbc azure-core + + python - < AccessToken: } -def get_pyodbc_attrs_before(credentials: FabricCredentials) -> Dict: +def get_pyodbc_attrs_before_credentials(credentials: FabricCredentials) -> Dict: """ Get the pyodbc attrs before. @@ -220,6 +220,36 @@ def get_pyodbc_attrs_before(credentials: FabricCredentials) -> Dict: return attrs_before +def get_pyodbc_attrs_before_accesstoken(accessToken: str) -> Dict: + """ + Get the pyodbc attrs before. + + Parameters + ---------- + credentials : Access Token for Integration Tests + Credentials. + + Returns + ------- + out : Dict + The pyodbc attrs before. + + Source + ------ + Authentication for SQL server with an access token: + https://docs.microsoft.com/en-us/sql/connect/odbc/using-azure-active-directory?view=sql-server-ver15#authenticating-with-an-access-token + """ + + access_token_utf16 = accessToken.encode("utf-16-le") + token_struct = struct.pack( + f" str: """ Convert a boolean to a connection string argument. @@ -323,7 +353,7 @@ def open(cls, connection: Connection) -> Connection: con_str.append(f"Database={credentials.database}") - #Enabling trace flag + # Enabling trace flag if credentials.trace_flag: con_str.append("SQL_ATTR_TRACE=SQL_OPT_TRACE_ON") else: @@ -395,7 +425,11 @@ def open(cls, connection: Connection) -> Connection: def connect(): logger.debug(f"Using connection string: {con_str_display}") - attrs_before = get_pyodbc_attrs_before(credentials) + if credentials.authentication == "ActiveDirectoryAccessToken": + attrs_before = get_pyodbc_attrs_before_accesstoken(credentials.access_token) + else: + attrs_before = get_pyodbc_attrs_before_credentials(credentials) + handle = pyodbc.connect( con_str_concat, attrs_before=attrs_before, diff --git a/dbt/adapters/fabric/fabric_credentials.py b/dbt/adapters/fabric/fabric_credentials.py index a824fac..138e3bd 100644 --- a/dbt/adapters/fabric/fabric_credentials.py +++ b/dbt/adapters/fabric/fabric_credentials.py @@ -17,6 +17,7 @@ class FabricCredentials(Credentials): tenant_id: Optional[str] = None client_id: Optional[str] = None client_secret: Optional[str] = None + access_token: Optional[str] = None authentication: Optional[str] = "ActiveDirectoryServicePrincipal" encrypt: Optional[bool] = True # default value in MS ODBC Driver 18 as well trust_cert: Optional[bool] = False # default value in MS ODBC Driver 18 as well diff --git a/tests/conftest.py b/tests/conftest.py index 3e60ce0..72eb2d6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -27,6 +27,8 @@ def dbt_profile_target(request: FixtureRequest, dbt_profile_target_update): target = _profile_ci_azure_environment() elif profile == "user_azure": target = _profile_user_azure() + elif profile == "integration_tests": + target = _profile_integration_tests() else: raise ValueError(f"Unknown profile: {profile}") @@ -55,7 +57,7 @@ def _profile_ci_azure_base(): "database": os.getenv("DBT_AZURESQL_DB"), "encrypt": True, "trust_cert": True, - "trace_flag":False, + "trace_flag": False, }, } @@ -104,6 +106,21 @@ def _profile_user_azure(): return profile +def _profile_integration_tests(): + profile = { + **_all_profiles_base(), + **{ + "host": os.getenv("FABRIC_TEST_HOST"), + "authentication": os.getenv("FABRIC_TEST_AUTH", "ActiveDirectoryAccessToken"), + "encrypt": True, + "trust_cert": True, + "database": os.getenv("FABRIC_TEST_DBNAME"), + "access_token": os.getenv("FABRIC_INTEGRATION_TESTS_TOKEN"), + }, + } + return profile + + @pytest.fixture(autouse=True) def skip_by_profile_type(request: FixtureRequest): profile_type = request.config.getoption("--profile")