rusackas commented on code in PR #41421:
URL: https://github.com/apache/superset/pull/41421#discussion_r3476836484


##########
superset/db_engine_specs/databricks.py:
##########
@@ -474,6 +576,74 @@ class 
DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec):
     supports_dynamic_catalog = True
     supports_cross_catalog_queries = True
 
+    # OAuth 2.0 support
+    supports_oauth2 = True
+    oauth2_exception = OAuth2RedirectError
+    oauth2_scope = "sql"
+
+    # OAuth2 endpoints are determined dynamically based on cloud provider
+    oauth2_authorization_request_uri = ""  # Set dynamically
+    oauth2_token_request_uri = ""  # Set dynamically
+
+    @classmethod
+    def get_oauth2_authorization_uri(
+        cls,
+        config: "OAuth2ClientConfig",
+        state: "OAuth2State",
+        code_verifier: str | None = None,
+    ) -> str:
+        """
+        Return URI for initial OAuth2 request with dynamic endpoint detection.
+
+        A fully-resolved `authorization_request_uri` from 
`DATABASE_OAUTH2_CLIENTS`
+        is preserved; only fall back to the auto-detected, account-resolved
+        endpoint when none is configured.
+        """
+        if not config.get("authorization_request_uri"):
+            from superset import db
+            from superset.models.core import Database
+
+            # Get the database to detect cloud provider
+            database_id = state["database_id"]
+            if database := db.session.get(Database, database_id):
+                provider = cls._detect_cloud_provider(database)
+                from typing import cast
+
+                config = cast(
+                    "OAuth2ClientConfig",
+                    dict(config)
+                    | {
+                        "authorization_request_uri": 
cls._resolve_oauth2_endpoint(
+                            database, provider, "authorization_request_uri"
+                        )
+                    },
+                )
+
+        return super().get_oauth2_authorization_uri(config, state, 
code_verifier)
+
+    @classmethod
+    def get_oauth2_token(
+        cls,
+        config: "OAuth2ClientConfig",
+        code: str,
+        code_verifier: str | None = None,
+    ) -> "OAuth2TokenResponse":
+        """
+        Exchange authorization code for refresh/access tokens.
+
+        The token request URI is resolved when the OAuth2 config is built (see
+        `get_oauth2_config`) and already targets the correct cloud provider and
+        account. There is no database context here to auto-detect it, so fail
+        fast rather than POST to an unresolved endpoint when it is missing.
+        """
+        if not config.get("token_request_uri"):
+            raise OAuth2Error(
+                "Databricks OAuth2 token endpoint is not configured: provide a 
"
+                "fully-resolved `token_request_uri` in 
DATABASE_OAUTH2_CLIENTS."
+            )
+
+        return super().get_oauth2_token(config, code, code_verifier)

Review Comment:
   You're right that this is asymmetric, the auth URI auto-resolves but 
`get_oauth2_token` has no database context to do the same, so a config without 
`token_request_uri` dies at exchange. There isn't a clean way to thread the 
provider/account in here without changing the base signature, so I'm inclined 
to require `token_request_uri` for the auto-detect path and tighten the docs to 
say so, rather than half-promise it. Will sort that out before this lands.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to