bhavishya123 commented on code in PR #36850:
URL: https://github.com/apache/superset/pull/36850#discussion_r2650472664
##########
superset/db_engine_specs/databricks.py:
##########
@@ -218,43 +227,221 @@ class DatabricksPythonConnectorPropertiesType(TypedDict):
class DatabricksHiveEngineSpec(HiveEngineSpec):
- engine_name = "Databricks Interactive Cluster"
+ """
+ Engine spec for Databricks Interactive Cluster using PyHive.
+ Note: This inherits from HiveEngineSpec, not DatabricksBaseEngineSpec,
+ to preserve Hive-specific functionality.
+ """
+ engine_name = "Databricks Interactive Cluster"
engine = "databricks"
drivers = {"pyhive": "Hive driver for Interactive Cluster"}
default_driver = "pyhive"
_show_functions_column = "function"
-
_time_grain_expressions = time_grain_expressions
+ # OAuth 2.0 support - replicate from DatabricksBaseEngineSpec since we
inherit from HiveEngineSpec
+ allows_user_impersonation = True
+ supports_oauth2 = True
+ oauth2_scope = "all-apis offline_access"
+ oauth2_exception = RequestError
+
+ @classmethod
+ def is_oauth2_enabled(cls) -> bool:
+ """Check if OAuth2 is enabled for this engine."""
+ return (
+ cls.supports_oauth2
+ and cls.engine_name in app.config.get("DATABASE_OAUTH2_CLIENTS",
{})
+ )
+
+ @classmethod
+ def needs_oauth2(cls, ex: Exception) -> bool:
+ """Check if the exception indicates that OAuth2 authentication is
needed."""
+ if not (g and hasattr(g, "user")):
+ return False
+
+ error_msg = str(ex).lower()
+ auth_error_patterns = [
+ "credential was not sent",
+ "unsupported type for this api",
+ "no valid authentication settings",
+ "authentication failed",
+ "invalid access token",
+ "token expired",
+ "unauthorized",
+ "401",
+ "403",
+ ]
+ return any(pattern in error_msg for pattern in auth_error_patterns)
+
+ @classmethod
+ def impersonate_user(
+ cls,
+ database: "Database",
+ username: str | None,
+ user_token: str | None,
+ url: URL,
+ engine_kwargs: dict[str, Any],
+ ) -> tuple[URL, dict[str, Any]]:
+ """Update the connection to use OAuth2 access token."""
+ if user_token:
+ url = url.set(password=user_token)
+ return url, engine_kwargs
+
+ @staticmethod
+ def update_params_from_encrypted_extra(
+ database: "Database", params: dict[str, Any]
+ ) -> None:
+ """Remove OAuth2 client info from engine params."""
+ HiveEngineSpec.update_params_from_encrypted_extra(database, params)
+ if "oauth2_client_info" in params:
+ del params["oauth2_client_info"]
+
Review Comment:
The OAuth2 implementation in DatabricksHiveEngineSpec is intentional, not
duplication. This class must inherit from HiveEngineSpec to preserve
Hive-specific behavior, so it cannot inherit from DatabricksBaseEngineSpec.
This pattern keeps the inheritance hierarchy clear and follows existing
Superset conventions for engine specs with different parent classes
--
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]
