docstring

Author: nicornk

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

@@ -326,7 +326,7 @@ def query_foundry_sql(
         arrow_compression_codec: ArrowCompressionCodec = ...,
         timeout: int = ...,
         experimental_use_trino: bool = ...,
-    ) -> pd.core.frame.DataFrame: ...
+    ) -> pd.DataFrame: ...
 
     @overload
     def query_foundry_sql(
@@ -374,7 +374,7 @@ def query_foundry_sql(
         arrow_compression_codec: ArrowCompressionCodec = ...,
         timeout: int = ...,
         experimental_use_trino: bool = ...,
-    ) -> pd.core.frame.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
+    ) -> pd.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
 
     def query_foundry_sql(
         self,
@@ -385,7 +385,7 @@ def query_foundry_sql(
         arrow_compression_codec: ArrowCompressionCodec = "NONE",
         timeout: int = 600,
         experimental_use_trino: bool = False,
-    ) -> pd.core.frame.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame:
+    ) -> pd.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame:
         """Queries the Foundry SQL server using the V2 API.
 
         Uses Arrow IPC to communicate with the Foundry SQL Server Endpoint.
@@ -510,7 +510,7 @@ def _extract_ticket(self, response_json: dict[str, Any]) -> dict[str, Any]:
             response_json: Success response JSON from status API
 
         Returns:
-            List of tickets for fetching results
+            Ticket dict with id, tickets list, and type. Example: {"id": 0, "tickets": [...], "type": "furnace"}
 
         Raises:
             KeyError: If the response JSON doesn't contain the expected structure

tests/integration/clients/test_foundry_sql_server.py

@@ -168,6 +168,7 @@ def test_v2_disable_arrow_compression():
     """Test V2 client with arrow compression disabled."""
     result = TEST_SINGLETON.ctx.foundry_sql_server_v2.query_foundry_sql(
         query=f"SELECT * FROM `{TEST_SINGLETON.iris_new.rid}` LIMIT 5",
+        arrow_compression_codec="NONE",
     )
     assert result.shape[0] == 5