|
| 1 | +import os |
| 2 | +import threading |
1 | 3 | import time |
2 | 4 | from functools import cached_property |
3 | 5 | from typing import Any, Dict |
|
13 | 15 |
|
14 | 16 | SUBMISSION_LANGUAGE = "python" |
15 | 17 |
|
| 18 | +# Minimum remaining token lifetime before refreshing (seconds). |
| 19 | +_TOKEN_REFRESH_MARGIN_SECONDS = 120 |
| 20 | + |
| 21 | +_spark_connect_env_lock = threading.Lock() |
| 22 | +_spark_connect_env_set = False |
| 23 | + |
| 24 | + |
| 25 | +def _ensure_spark_connect_env() -> None: |
| 26 | + """Set SPARK_CONNECT_MODE_ENABLED=1 once for the process.""" |
| 27 | + global _spark_connect_env_set |
| 28 | + if _spark_connect_env_set: |
| 29 | + return |
| 30 | + with _spark_connect_env_lock: |
| 31 | + if not _spark_connect_env_set: |
| 32 | + os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" |
| 33 | + _spark_connect_env_set = True |
| 34 | + |
| 35 | + |
| 36 | +def _create_athena_channel_builder( |
| 37 | + athena_client: Any, |
| 38 | + session_id: str, |
| 39 | + endpoint_url: str, |
| 40 | + initial_auth_token: str | None = None, |
| 41 | + initial_token_expiry: Any = None, |
| 42 | +) -> Any: |
| 43 | + """Create a ChannelBuilder subclass that auto-refreshes the Athena AuthToken. |
| 44 | +
|
| 45 | + The AuthToken from GetSessionEndpoint expires after ~29 minutes. |
| 46 | + This builder calls GetSessionEndpoint to obtain a fresh token |
| 47 | + whenever the current token is about to expire, so that long-running |
| 48 | + Spark Connect jobs are not interrupted by PERMISSION_DENIED errors. |
| 49 | +
|
| 50 | + Returns a ChannelBuilder subclass instance (import deferred to avoid |
| 51 | + top-level pyspark dependency). |
| 52 | + """ |
| 53 | + from datetime import datetime, timezone |
| 54 | + |
| 55 | + from pyspark.sql.connect.client.core import ChannelBuilder |
| 56 | + |
| 57 | + class AthenaChannelBuilder(ChannelBuilder): |
| 58 | + def __init__(self, client: Any, sid: str, url: str, auth_token: str | None, token_expiry: Any): |
| 59 | + sc_url = url.replace("https://", "sc://", 1) + ":443/;use_ssl=true" |
| 60 | + super().__init__(sc_url) |
| 61 | + self._athena_client = client |
| 62 | + self._athena_session_id = sid |
| 63 | + self._auth_token = auth_token |
| 64 | + self._token_expiry = token_expiry |
| 65 | + |
| 66 | + def _refresh_token_if_needed(self) -> None: |
| 67 | + if self._auth_token and self._token_expiry: |
| 68 | + remaining = (self._token_expiry - datetime.now(timezone.utc)).total_seconds() |
| 69 | + if remaining > _TOKEN_REFRESH_MARGIN_SECONDS: |
| 70 | + return |
| 71 | + LOGGER.debug(f"AuthToken expiring in {remaining:.0f}s, refreshing") |
| 72 | + |
| 73 | + response = self._athena_client.get_session_endpoint( |
| 74 | + SessionId=self._athena_session_id |
| 75 | + ) |
| 76 | + auth_token = response.get("AuthToken") |
| 77 | + if not auth_token: |
| 78 | + raise DbtRuntimeError( |
| 79 | + f"GetSessionEndpoint returned no AuthToken for session {self._athena_session_id}" |
| 80 | + ) |
| 81 | + self._auth_token = auth_token |
| 82 | + self._token_expiry = response.get("AuthTokenExpirationTime") |
| 83 | + |
| 84 | + def metadata(self): |
| 85 | + self._refresh_token_if_needed() |
| 86 | + base = [(k, v) for k, v in super().metadata() if k != "x-aws-proxy-auth"] |
| 87 | + base.append(("x-aws-proxy-auth", self._auth_token)) |
| 88 | + return base |
| 89 | + |
| 90 | + return AthenaChannelBuilder(athena_client, session_id, endpoint_url, initial_auth_token, initial_token_expiry) |
| 91 | + |
16 | 92 |
|
17 | 93 | class AthenaPythonJobHelper(PythonJobHelper): |
18 | 94 | """ |
@@ -42,6 +118,7 @@ def __init__(self, parsed_model: Dict[Any, Any], credentials: AthenaCredentials) |
42 | 118 | self.polling_interval, |
43 | 119 | self.engine_config, |
44 | 120 | self.relation_name, |
| 121 | + spark_managed_logging=self.config.config.get("spark_managed_logging", False), |
45 | 122 | ) |
46 | 123 |
|
47 | 124 | @cached_property |
@@ -107,22 +184,121 @@ def submit(self, compiled_code: str) -> Any: |
107 | 184 | """ |
108 | 185 | Submit a calculation to Athena. |
109 | 186 |
|
110 | | - This function submits a calculation to Athena for execution using the provided compiled code. |
111 | | - It starts a calculation execution with the current session ID and the compiled code as the code block. |
112 | | - The function then polls until the calculation execution is completed, and retrieves the result. |
113 | | - If the execution is successful and completed, the result S3 URI is returned. Otherwise, a DbtRuntimeError |
114 | | - is raised with the execution status. |
| 187 | + For PySpark engine version 3, uses the Calculations API |
| 188 | + (StartCalculationExecution). |
| 189 | + For Apache Spark version 3.5+, uses Spark Connect via |
| 190 | + GetSessionEndpoint. |
| 191 | + """ |
| 192 | + if self.config.config.get("spark_engine_version") == "3.5": |
| 193 | + return self._submit_spark_connect(compiled_code) |
| 194 | + return self._submit_calculation_api(compiled_code) |
115 | 195 |
|
116 | | - Args: |
117 | | - compiled_code (str): The compiled code to submit for execution. |
| 196 | + def _wait_for_endpoint(self) -> Dict[str, Any]: |
| 197 | + """Poll until the session endpoint is ready and return the full response.""" |
| 198 | + import random |
118 | 199 |
|
119 | | - Returns: |
120 | | - dict: The result S3 URI if the execution is successful and completed. |
| 200 | + polling_interval = self.polling_interval |
| 201 | + timer: float = 0 |
| 202 | + throttle_backoff: float = 0 |
| 203 | + while True: |
| 204 | + try: |
| 205 | + response = self.athena_client.get_session_endpoint( |
| 206 | + SessionId=self.session_id |
| 207 | + ) |
| 208 | + endpoint_url = response.get("EndpointUrl") |
| 209 | + if endpoint_url: |
| 210 | + if not response.get("AuthToken"): |
| 211 | + raise DbtRuntimeError( |
| 212 | + f"GetSessionEndpoint returned no AuthToken for session {self.session_id}" |
| 213 | + ) |
| 214 | + return response |
| 215 | + throttle_backoff = 0 |
| 216 | + except botocore.exceptions.ClientError as e: |
| 217 | + error_code = e.response.get("Error", {}).get("Code", "") |
| 218 | + if error_code == "ThrottlingException": |
| 219 | + throttle_backoff = min((throttle_backoff or 1) * 2, 30) + random.uniform(0, 1) |
| 220 | + LOGGER.warning( |
| 221 | + f"Session {self.session_id} endpoint throttled, backing off {throttle_backoff:.1f}s" |
| 222 | + ) |
| 223 | + else: |
| 224 | + throttle_backoff = 0 |
| 225 | + LOGGER.debug(f"Waiting for session {self.session_id} endpoint: {e}") |
| 226 | + if timer >= self.timeout: |
| 227 | + raise DbtRuntimeError( |
| 228 | + f"Session {self.session_id} endpoint did not become available within {self.timeout}s" |
| 229 | + ) |
| 230 | + sleep_time = throttle_backoff if throttle_backoff else polling_interval |
| 231 | + time.sleep(sleep_time) |
| 232 | + timer += sleep_time |
121 | 233 |
|
122 | | - Raises: |
123 | | - DbtRuntimeError: If the execution ends in a state other than "COMPLETED". |
| 234 | + def _submit_spark_connect(self, compiled_code: str) -> Any: |
| 235 | + """Submit code via Spark Connect (Apache Spark version 3.5+). |
124 | 236 |
|
| 237 | + Uses AthenaChannelBuilder to auto-refresh the AuthToken before |
| 238 | + it expires (~29 min TTL), enabling long-running jobs. |
| 239 | + Enforces self.timeout as a hard execution time limit covering both |
| 240 | + endpoint wait and code execution. |
125 | 241 | """ |
| 242 | + if not compiled_code.strip(): |
| 243 | + return {} |
| 244 | + |
| 245 | + _ensure_spark_connect_env() |
| 246 | + |
| 247 | + spark = None |
| 248 | + timer = None |
| 249 | + timeout_event = threading.Event() |
| 250 | + start_time = time.monotonic() |
| 251 | + |
| 252 | + try: |
| 253 | + response = self._wait_for_endpoint() |
| 254 | + channel_builder = _create_athena_channel_builder( |
| 255 | + self.athena_client, |
| 256 | + self.session_id, |
| 257 | + response["EndpointUrl"], |
| 258 | + initial_auth_token=response.get("AuthToken"), |
| 259 | + initial_token_expiry=response.get("AuthTokenExpirationTime"), |
| 260 | + ) |
| 261 | + |
| 262 | + from pyspark.sql.connect.session import SparkSession as ConnectSparkSession |
| 263 | + |
| 264 | + spark = ConnectSparkSession.builder.channelBuilder(channel_builder).create() |
| 265 | + |
| 266 | + elapsed = time.monotonic() - start_time |
| 267 | + remaining = max(self.timeout - elapsed, 0) |
| 268 | + |
| 269 | + def _on_timeout(): |
| 270 | + timeout_event.set() |
| 271 | + LOGGER.warning( |
| 272 | + f"Model {self.relation_name} - Execution timed out after {self.timeout}s" |
| 273 | + ) |
| 274 | + spark.interruptAll() |
| 275 | + |
| 276 | + timer = threading.Timer(remaining, _on_timeout) |
| 277 | + timer.start() |
| 278 | + |
| 279 | + exec_globals = {"spark": spark} |
| 280 | + exec(compiled_code, exec_globals) |
| 281 | + except DbtRuntimeError: |
| 282 | + raise |
| 283 | + except Exception as e: |
| 284 | + if timeout_event.is_set(): |
| 285 | + raise DbtRuntimeError( |
| 286 | + f"Spark Connect execution timed out after {self.timeout} seconds." |
| 287 | + ) |
| 288 | + import traceback |
| 289 | + LOGGER.error(f"Spark Connect traceback:\n{traceback.format_exc()}") |
| 290 | + raise DbtRuntimeError(f"Spark Connect execution failed: {type(e).__name__}: {e}") from e |
| 291 | + finally: |
| 292 | + if timer is not None: |
| 293 | + timer.cancel() |
| 294 | + if spark is not None: |
| 295 | + spark.stop() |
| 296 | + self.spark_connection.set_spark_session_load(self.session_id, -1) |
| 297 | + |
| 298 | + return {} |
| 299 | + |
| 300 | + def _submit_calculation_api(self, compiled_code: str) -> Any: |
| 301 | + """Submit code via Calculations API (PySpark engine version 3).""" |
126 | 302 | # Seeing an empty calculation along with main python model code calculation is submitted for almost every model |
127 | 303 | # Also, if not returning the result json, we are getting green ERROR messages instead of OK messages. |
128 | 304 | # And with this handling, the run model code in target folder every model under run folder seems to be empty |
@@ -171,11 +347,14 @@ def submit(self, compiled_code: str) -> Any: |
171 | 347 | LOGGER.debug( |
172 | 348 | f"Model {self.relation_name} - Received execution status {execution_status}" |
173 | 349 | ) |
| 350 | + result = {} |
174 | 351 | if execution_status == "COMPLETED": |
175 | 352 | try: |
176 | | - result = self.athena_client.get_calculation_execution( |
| 353 | + execution = self.athena_client.get_calculation_execution( |
177 | 354 | CalculationExecutionId=calculation_execution_id |
178 | | - )["Result"] |
| 355 | + ) |
| 356 | + result = execution["Result"] |
| 357 | + result["Statistics"] = execution.get("Statistics", {}) |
179 | 358 | except Exception as e: |
180 | 359 | LOGGER.error(f"Unable to retrieve results: Got: {e}") |
181 | 360 | result = {} |
|
0 commit comments