bito-code-review[bot] commented on code in PR #36850:
URL: https://github.com/apache/superset/pull/36850#discussion_r2650474477
##########
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:
<!-- Bito Reply -->
Your explanation is accurate. DatabricksHiveEngineSpec inherits from
HiveEngineSpec to retain Hive-specific functionality, as noted in the class
docstring. The OAuth2 attributes and methods are replicated from
DatabricksBaseEngineSpec due to this inheritance constraint, following
Superset's pattern for engine specs with different parents.
--
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]
