feat: add support for azure access token authorization

duckdb · Sep 10, 2024 · 10b182f · 10b182f
1 parent 05a87bc
commit 10b182f
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 2 deletions.
diff --git a/.github/workflows/CloudTesting.yml b/.github/workflows/CloudTesting.yml
@@ -65,6 +65,14 @@ jobs:
         run: |
           python3 duckdb/scripts/run_tests_one_by_one.py ./build/release/test/unittest `pwd`/test/sql/cloud/*
 
+      - name: Test with Access Token in env vars
+        env:
+          AZURE_STORAGE_ACCOUNT: ${{secrets.AZURE_STORAGE_ACCOUNT}}
+        run: |
+          az login --service-principal -u ${{secrets.AZURE_CLIENT_ID}} -p ${{secrets.AZURE_CLIENT_SECRET}} --tenant ${{secrets.AZURE_TENANT_ID}}
+          export AZURE_ACCESS_TOKEN=`az account get-access-token --resource https://storage.azure.com --query accessToken --output tsv`        
+          python3 duckdb/scripts/run_tests_one_by_one.py ./build/release/test/unittest `pwd`/test/sql/cloud/*          
+
       - name: Test with SPN logged in in azure-cli
         env:
           AZURE_STORAGE_ACCOUNT: ${{secrets.AZURE_STORAGE_ACCOUNT}}

diff --git a/src/functions/delta_scan.cpp b/src/functions/delta_scan.cpp
@@ -283,8 +283,15 @@ static ffi::EngineBuilder* CreateBuilder(ClientContext &context, const string &p
         secret_reader.TryGetSecretKey("chain", chain);
 
         auto provider = kv_secret.GetProvider();
-
-        if (provider == "credential_chain") {
+        if (provider == "access_token") {
+            // Authentication option 0: https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html#variant.Token
+            string access_token;
+            secret_reader.TryGetSecretKey("access_token", access_token);
+            if (access_token.empty()) {
+                throw InvalidInputException("No access_token value not found in secret provider!");
+            }
+            ffi::set_builder_option(builder, KernelUtils::ToDeltaString("bearer_token"), KernelUtils::ToDeltaString(access_token));
+        } else if (provider == "credential_chain") {
             // Authentication option 1a: using the cli authentication
             if (chain.find("cli") != std::string::npos) {
                 ffi::set_builder_option(builder, KernelUtils::ToDeltaString("use_azure_cli"), KernelUtils::ToDeltaString("true"));

diff --git a/test/sql/cloud/azure/access_token_auth.test b/test/sql/cloud/azure/access_token_auth.test
@@ -0,0 +1,44 @@
+# name: test/sql/cloud/azure/access_token_auth.test
+# description: test access-token authentication
+# group: [azure]
+
+require azure
+
+require delta
+
+require-env AZURE_ACCESS_TOKEN
+
+require-env AZURE_STORAGE_ACCOUNT
+
+statement ok
+set allow_persistent_secrets=false
+
+statement error
+SELECT count(*) FROM delta_scan('azure://delta-testing-private/dat/all_primitive_types/delta');
+----
+Invalid Input Error: No valid Azure credentials found!
+
+statement ok
+CREATE SECRET az1 (
+    TYPE AZURE,
+    PROVIDER ACCESS_TOKEN,
+    ACCOUNT_NAME '${AZURE_STORAGE_ACCOUNT}'
+)
+
+statement error
+SELECT count(*) FROM delta_scan('azure://delta-testing-private/dat/all_primitive_types/delta');
+----
+Invalid Input Error: No access_token value not found in secret provider!
+
+statement ok
+CREATE OR REPLACE SECRET az1 (
+    TYPE AZURE,
+    PROVIDER ACCESS_TOKEN,
+    ACCESS_TOKEN '${AZURE_ACCESS_TOKEN}',
+    ACCOUNT_NAME '${AZURE_STORAGE_ACCOUNT}'
+)
+
+query I
+SELECT count(*) FROM delta_scan('azure://delta-testing-private/dat/all_primitive_types/delta');
+----
+5