From de75bac50805eede2f59024cc249b1a7bb5f7f7b Mon Sep 17 00:00:00 2001
From: Pradeep Srikakolapu <prdpsvs@gmail.com>
Date: Mon, 15 Apr 2024 10:52:45 -0700
Subject: [PATCH] Adding Synapse Spark Authentication Option

---
 dbt/adapters/fabric/__version__.py            |  2 +-
 .../fabric/fabric_connection_manager.py       | 25 +++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/dbt/adapters/fabric/__version__.py b/dbt/adapters/fabric/__version__.py
index 72126ce..ba22724 100644
--- a/dbt/adapters/fabric/__version__.py
+++ b/dbt/adapters/fabric/__version__.py
@@ -1 +1 @@
-version = "1.8.1"
+version = "1.8.2"
diff --git a/dbt/adapters/fabric/fabric_connection_manager.py b/dbt/adapters/fabric/fabric_connection_manager.py
index 351317c..9719d84 100644
--- a/dbt/adapters/fabric/fabric_connection_manager.py
+++ b/dbt/adapters/fabric/fabric_connection_manager.py
@@ -23,6 +23,7 @@
 from dbt.adapters.fabric.fabric_credentials import FabricCredentials
 
 AZURE_CREDENTIAL_SCOPE = "https://database.windows.net//.default"
+SYNAPSE_SPARK_CREDENTIAL_SCOPE = "DW"
 _TOKEN: Optional[AccessToken] = None
 AZURE_AUTH_FUNCTION_TYPE = Callable[[FabricCredentials], AccessToken]
 
@@ -87,6 +88,29 @@ def convert_access_token_to_mswindows_byte_string(token: AccessToken) -> bytes:
     return convert_bytes_to_mswindows_byte_string(value)
 
 
+def get_synapse_spark_access_token(credentials: FabricCredentials) -> AccessToken:
+    """
+    Get an Azure access token by using mspsarkutils
+    Parameters
+    -----------
+    credentials: FabricCredentials
+        Credentials.
+    Returns
+    -------
+    out : AccessToken
+        The access token.
+    """
+    from notebookutils import mssparkutils
+
+    aad_token = mssparkutils.credentials.getToken(SYNAPSE_SPARK_CREDENTIAL_SCOPE)
+    expires_on = int(time.time() + 4500.0)
+    token = AccessToken(
+        token=aad_token,
+        expires_on=expires_on,
+    )
+    return token
+
+
 def get_cli_access_token(credentials: FabricCredentials) -> AccessToken:
     """
     Get an Azure access token using the CLI credentials
@@ -152,6 +176,7 @@ def get_environment_access_token(credentials: FabricCredentials) -> AccessToken:
     "cli": get_cli_access_token,
     "auto": get_auto_access_token,
     "environment": get_environment_access_token,
+    "synapsespark": get_synapse_spark_access_token,
 }