Skip to content

Commit d02d515

Browse files
committed
Fix response format
1 parent 21ac0b6 commit d02d515

1 file changed

Lines changed: 7 additions & 6 deletions

File tree

libs/foundry-dev-tools/src/foundry_dev_tools/clients/foundry_sql_server.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from __future__ import annotations
44

5+
import logging
56
import time
67
import warnings
78
from typing import TYPE_CHECKING, Any, Literal, overload
@@ -22,6 +23,8 @@
2223
import pyspark
2324
import requests
2425

26+
LOGGER = logging.getLogger(__name__)
27+
2528

2629
class FoundrySqlServerClient(APIClient):
2730
"""FoundrySqlServerClient class that implements methods from the 'foundry-sql-server' API."""
@@ -433,12 +436,10 @@ def _extract_query_identifier(self, response_json: dict[str, Any]) -> dict[str,
433436
Query identifier dict
434437
435438
"""
436-
if response_json.get("type") == "pending":
437-
return response_json["pending"]["query"]
438-
if response_json.get("type") == "success":
439-
return response_json["success"]["query"]
440-
msg = f"Unexpected response type: {response_json.get('type')}"
441-
raise ValueError(msg)
439+
if response_json["type"] == "triggered" and "plan" in response_json["triggered"]:
440+
plan = response_json["triggered"]["plan"]
441+
LOGGER.debug("plan %s", plan)
442+
return response_json[response_json["type"]]["query"]
442443

443444
def _extract_tickets(self, response_json: dict[str, Any]) -> list[str]:
444445
"""Extract tickets from success response.

0 commit comments

Comments
 (0)