Support for connecting via Service Principals

[kartikay-bagla]

Currently the only authentication method supported is via personal access tokens:

f"databricks://token:{access_token}@{host}?http_path={http_path}&catalog={catalog}&schema={schema}"

Technically, we can generate tokens for service principals by calling the Workspace OAuth API. But that returns a short-lived access token, and when that token expires, we would need to create a new engine with the updated token.

The Databricks SQL Connector for Python already supports M2M authentication which removes the overhead of managing short-lived access tokens.

Would it be possible to have connector itself manage refreshing the access tokens? The connection string would then possibly become:

f"databricks://{service_principal_uuid}:{service_principal_secret}@{host}?http_path={http_path}&catalog={catalog}&schema={schema}"

[amirhessam88]

You just need to add that functionality via sqlalchemy to inject fresh token

from sqlalchemy import event

engine = create_engine("postgresql+psycopg2://user:pass@hostname/dbname")


@event.listens_for(engine, "do_connect")
def receive_do_connect(dialect, conn_rec, cargs, cparams):

    # add a function to grab the fresh token via API
    new_token = get_new_token()
    cparams["access_token"] = new token

    return engine

https://docs.sqlalchemy.org/en/20/core/engines.html#fully-replacing-the-dbapi-connect-function

[dhirschfeld]

Thanks for the work-around Amirhessam - that's a good escape-hatch, however I'd still like for service principals to be supported natively by databricks-sqlalchemy so that I don't have to maintain awful hacks in my own code.

[jprakash-db]

Hi everyone! Our goal is to keep the databricks-sqlalchemy very lean and its main goal is just to serve as a wrapper around sqlalchemy, so we don't have plans to introduce any other logic within databricks-sqlalchemy.
So to utilize any auth related functionality you must use the sql connector.
In the databricks-sqlalchemy there is a argument called connect_args which can be used to pass anything to the sql connector.
So for any auth, just pass the required arguments in connect_args as you would pass to the sql connector.

Ex if you want to use the azure m2m

from sqlalchemy import create_engine, text
import os
from databricks.sql.auth.auth import AuthType

# Azure Oauth
client_id = "Azure_Client_ID"
client_secret = "Azure_Client_Secret"
host_name="Host_Name"
http_path="Http_Path"
auth_type=AuthType.AZURE_SP_M2M.value
catalog = "Catalog"
access_token = "dummy_token"
schema = "schema"

engine = create_engine(
    f"databricks://token:{access_token}@{host_name}?http_path={http_path}&catalog={catalog}&schema={schema}",
    connect_args={
     "auth_type":auth_type,
     "azure_client_id":client_id,
     "azure_client_secret":client_secret,
    }
)
# 1. Get all catalogs (databases)
with engine.connect() as conn:
    result = conn.execute(text("SELECT 1"))
    query_result = [row[0] for row in result]
    print("Result:", query_result)