bito-code-review[bot] commented on code in PR #36850:
URL: https://github.com/apache/superset/pull/36850#discussion_r2650483630
##########
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 -->
Yes, apply the suggested try-except guard in is_oauth2_enabled to prevent
RuntimeError when accessing app.config without an application context. This
fixes the potential bug flagged in the comment.
**superset/db_engine_specs/databricks.py**
```
@classmethod
def is_oauth2_enabled(cls) -> bool:
"""Check if OAuth2 is enabled for this engine."""
try:
clients = app.config.get("DATABASE_OAUTH2_CLIENTS", {})
except RuntimeError:
# No app context available; treat OAuth2 as disabled in this
context
return False
return cls.supports_oauth2 and cls.engine_name in clients
```
--
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]
