add unit tests

Author: nicornk

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

@@ -14,7 +14,14 @@
     FoundrySqlQueryFailedError,
     FoundrySqlSerializationFormatNotImplementedError,
 )
-from foundry_dev_tools.utils.api_types import ArrowCompressionCodec, Ref, SqlDialect, SQLReturnType, assert_in_literal
+from foundry_dev_tools.utils.api_types import (
+    ArrowCompressionCodec,
+    FurnaceSqlDialect,
+    Ref,
+    SqlDialect,
+    SQLReturnType,
+    assert_in_literal,
+)
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -318,9 +325,10 @@ def query_foundry_sql(
         query: str,
         return_type: Literal["pandas"],
         branch: Ref = ...,
-        sql_dialect: SqlDialect = ...,
+        sql_dialect: FurnaceSqlDialect = ...,
         arrow_compression_codec: ArrowCompressionCodec = ...,
         timeout: int = ...,
+        experimental_use_trino: bool = ...,
     ) -> pd.core.frame.DataFrame: ...
 
     @overload
@@ -329,9 +337,10 @@ def query_foundry_sql(
         query: str,
         return_type: Literal["polars"],
         branch: Ref = ...,
-        sql_dialect: SqlDialect = ...,
+        sql_dialect: FurnaceSqlDialect = ...,
         arrow_compression_codec: ArrowCompressionCodec = ...,
         timeout: int = ...,
+        experimental_use_trino: bool = ...,
     ) -> pl.DataFrame: ...
 
     @overload
@@ -340,9 +349,10 @@ def query_foundry_sql(
         query: str,
         return_type: Literal["spark"],
         branch: Ref = ...,
-        sql_dialect: SqlDialect = ...,
+        sql_dialect: FurnaceSqlDialect = ...,
         arrow_compression_codec: ArrowCompressionCodec = ...,
         timeout: int = ...,
+        experimental_use_trino: bool = ...,
     ) -> pyspark.sql.DataFrame: ...
 
     @overload
@@ -351,9 +361,10 @@ def query_foundry_sql(
         query: str,
         return_type: Literal["arrow"],
         branch: Ref = ...,
-        sql_dialect: SqlDialect = ...,
+        sql_dialect: FurnaceSqlDialect = ...,
         arrow_compression_codec: ArrowCompressionCodec = ...,
         timeout: int = ...,
+        experimental_use_trino: bool = ...,
     ) -> pa.Table: ...
 
     @overload
@@ -362,19 +373,21 @@ def query_foundry_sql(
         query: str,
         return_type: SQLReturnType = ...,
         branch: Ref = ...,
-        sql_dialect: SqlDialect = ...,
+        sql_dialect: FurnaceSqlDialect = ...,
         arrow_compression_codec: ArrowCompressionCodec = ...,
         timeout: int = ...,
+        experimental_use_trino: bool = ...,
     ) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame: ...
 
     def query_foundry_sql(
         self,
         query: str,
         return_type: SQLReturnType = "pandas",
         branch: Ref = "master",
-        sql_dialect: SqlDialect = "SPARK",
+        sql_dialect: FurnaceSqlDialect = "SPARK",
         arrow_compression_codec: ArrowCompressionCodec = "NONE",
         timeout: int = 600,
+        experimental_use_trino: bool = False,
     ) -> tuple[dict, list[list]] | pd.core.frame.DataFrame | pl.DataFrame | pa.Table | pyspark.sql.DataFrame:
         """Queries the Foundry SQL server using the V2 API.
 
@@ -389,9 +402,10 @@ def query_foundry_sql(
             query: The SQL Query
             return_type: See :py:class:foundry_dev_tools.foundry_api_client.SQLReturnType
             branch: The dataset branch to query
-            sql_dialect: The SQL dialect to use
+            sql_dialect: The SQL dialect to use (only SPARK is supported for V2)
             arrow_compression_codec: Arrow compression codec (NONE, LZ4, ZSTD)
             timeout: Query timeout in seconds
+            experimental_use_trino: If True, modifies the query to use Trino backend by adding /*+ backend(trino) */ hint
 
         Returns:
             :external+pandas:py:class:`~pandas.DataFrame` | :external+polars:py:class:`~polars.DataFrame` | :external+pyarrow:py:class:`~pyarrow.Table` | :external+spark:py:class:`~pyspark.sql.DataFrame`:
@@ -403,6 +417,11 @@ def query_foundry_sql(
             FoundrySqlQueryClientTimedOutError: If the query times out
 
         """  # noqa: E501
+        assert_in_literal(sql_dialect, FurnaceSqlDialect, "sql_dialect")
+
+        if experimental_use_trino:
+            query = query.replace("SELECT ", "SELECT /*+ backend(trino) */ ", 1)
+
         response_json = self.api_query(
             query=query, dialect=sql_dialect, branch=branch, arrow_compression_codec=arrow_compression_codec
         ).json()
@@ -507,7 +526,7 @@ def read_stream_results_arrow(self, ticket: dict[str, Any]) -> pa.ipc.RecordBatc
     def api_query(
         self,
         query: str,
-        dialect: SqlDialect,
+        dialect: FurnaceSqlDialect,
         branch: Ref,
         arrow_compression_codec: ArrowCompressionCodec = "NONE",
         **kwargs,
@@ -516,7 +535,7 @@ def api_query(
 
         Args:
             query: The SQL query string
-            dialect: The SQL dialect to use
+            dialect: The SQL dialect to use (only SPARK is supported)
             branch: The dataset branch to query
             arrow_compression_codec: Arrow compression codec (NONE, LZ4, ZSTD)
             **kwargs: gets passed to :py:meth:`APIClient.api_request`

libs/foundry-dev-tools/src/foundry_dev_tools/utils/api_types.py

@@ -95,6 +95,9 @@ def assert_in_literal(option, literal, variable_name) -> None:  # noqa: ANN001
 SqlDialect = Literal["ANSI", "SPARK"]
 """The SQL Dialect for Foundry SQL queries."""
 
+FurnaceSqlDialect = Literal["SPARK"]
+"""The SQL Dialect for Furnace SQL queries (V2 API). Only SPARK is supported."""
+
 ArrowCompressionCodec = Literal["NONE", "LZ4", "ZSTD"]
 """The Arrow compression codec for Foundry SQL queries."""
 

tests/integration/clients/test_foundry_sql_server.py

@@ -73,6 +73,25 @@ def test_legacy_fallback(mocker):
     query_foundry_sql_legacy_spy.assert_called()
 
 
+def test_v1_ansi_sql_dialect():
+    """Test V1 client with ANSI SQL dialect (uses double quotes instead of backticks)."""
+    # Test basic query with ANSI dialect - note the use of double quotes instead of backticks
+    result = TEST_SINGLETON.ctx.foundry_sql_server.query_foundry_sql(
+        query=f'SELECT sepal_width, sepal_length FROM "{TEST_SINGLETON.iris_new.rid}" LIMIT 5',
+        sql_dialect="ANSI",
+    )
+    assert result.shape[0] == 5
+    assert result.shape[1] == 2
+
+    # Test with aggregation using ANSI dialect
+    result_agg = TEST_SINGLETON.ctx.foundry_sql_server.query_foundry_sql(
+        query=f'SELECT COUNT(*) as cnt FROM "{TEST_SINGLETON.iris_new.rid}"',
+        sql_dialect="ANSI",
+    )
+    assert result_agg.shape[0] == 1
+    assert "cnt" in result_agg.columns
+
+
 # V2 Client Tests
 
 
@@ -176,3 +195,73 @@ def test_v2_polars_return_type():
     assert isinstance(polars_df, pl.DataFrame)
     assert polars_df.height == 2
     assert polars_df.width == 1
+
+
+def test_v2_polars_parquet():
+    polars_df = TEST_SINGLETON.ctx.foundry_sql_server_v2.query_foundry_sql(
+        f"SELECT sepal_length FROM `{TEST_SINGLETON.iris_parquet.rid}` LIMIT 2",
+        return_type="polars",
+    )
+    assert isinstance(polars_df, pl.DataFrame)
+    assert polars_df.height == 2
+    assert polars_df.width == 1
+
+    polars_df = TEST_SINGLETON.ctx.foundry_sql_server_v2.query_foundry_sql(
+        f"SELECT sepal_length FROM `{TEST_SINGLETON.iris_parquet.rid}` LIMIT 2",
+        return_type="polars",
+        experimental_use_trino=True,
+    )
+    assert isinstance(polars_df, pl.DataFrame)
+    assert polars_df.height == 2
+    assert polars_df.width == 1
+
+
+def test_v2_polars_parquet_hive_partitioning():
+    polars_df = TEST_SINGLETON.ctx.foundry_sql_server_v2.query_foundry_sql(
+        f"SELECT sepal_length FROM `{TEST_SINGLETON.iris_hive_partitioned.rid}` LIMIT 2",
+        return_type="polars",
+        experimental_use_trino=True,
+    )
+    assert isinstance(polars_df, pl.DataFrame)
+    assert polars_df.height == 2
+    assert polars_df.width == 1
+
+    polars_df = TEST_SINGLETON.ctx.foundry_sql_server_v2.query_foundry_sql(
+        f"SELECT sepal_length FROM `{TEST_SINGLETON.iris_hive_partitioned.rid}` LIMIT 2", return_type="polars"
+    )
+    assert isinstance(polars_df, pl.DataFrame)
+    assert polars_df.height == 2
+    assert polars_df.width == 1
+
+
+def test_v2_arrow_compression_codecs():
+    """Test V2 client with different arrow compression codecs."""
+    # Test with LZ4 compression
+    result_lz4 = TEST_SINGLETON.ctx.foundry_sql_server_v2.query_foundry_sql(
+        query=f"SELECT * FROM `{TEST_SINGLETON.iris_new.rid}` LIMIT 10",
+        arrow_compression_codec="LZ4",
+    )
+    assert result_lz4.shape[0] == 10
+    assert result_lz4.shape[1] == 5
+
+    # Test with ZSTD compression
+    result_zstd = TEST_SINGLETON.ctx.foundry_sql_server_v2.query_foundry_sql(
+        query=f"SELECT * FROM `{TEST_SINGLETON.iris_new.rid}` LIMIT 10",
+        arrow_compression_codec="ZSTD",
+    )
+    assert result_zstd.shape[0] == 10
+    assert result_zstd.shape[1] == 5
+
+    # Test with NONE compression (default)
+    result_none = TEST_SINGLETON.ctx.foundry_sql_server_v2.query_foundry_sql(
+        query=f"SELECT * FROM `{TEST_SINGLETON.iris_new.rid}` LIMIT 10",
+        arrow_compression_codec="NONE",
+    )
+    assert result_none.shape[0] == 10
+    assert result_none.shape[1] == 5
+
+    # Verify all results have the same data
+    import pandas as pd
+
+    pd.testing.assert_frame_equal(result_lz4, result_zstd)
+    pd.testing.assert_frame_equal(result_lz4, result_none)

tests/unit/clients/test_foundry_sql_server.py

@@ -121,3 +121,106 @@ def test_exception_unknown_json(mocker, test_context_mock):
             timeout=0.001,
         )
     assert exception.value.error_message == ""
+
+
+def test_v2_experimental_use_trino(mocker, test_context_mock):
+    """Test that experimental_use_trino parameter modifies the query correctly."""
+    import pandas as pd
+
+    mocker.patch("time.sleep")  # we do not want to wait in tests
+
+    # Mock the arrow stream reader to return a simple pandas DataFrame
+    mock_arrow_reader = mocker.MagicMock()
+    mock_arrow_reader.read_pandas.return_value = pd.DataFrame({"col1": [1, 2, 3]})
+    mocker.patch.object(
+        test_context_mock.foundry_sql_server_v2,
+        "read_stream_results_arrow",
+        return_value=mock_arrow_reader,
+    )
+
+    # Mock the api_query endpoint (initial query execution)
+    query_matcher = mocker.MagicMock()
+    test_context_mock.mock_adapter.register_uri(
+        "POST",
+        build_api_url(TEST_HOST.url, "foundry-sql-server", "sql-endpoint/v1/queries/query"),
+        json={"type": "running", "running": {"queryHandle": {"queryId": "test-query-id-123", "type": "foundry"}}},
+        additional_matcher=query_matcher,
+    )
+
+    # Mock the api_status endpoint (poll for completion - returns ready immediately)
+    test_context_mock.mock_adapter.register_uri(
+        "POST",
+        build_api_url(TEST_HOST.url, "foundry-sql-server", "sql-endpoint/v1/queries/status"),
+        json={
+            "status": {
+                "type": "ready",
+                "ready": {"tickets": [{"tickets": ["eyJhbGc...mock-ticket-1", "eyJhbGc...mock-ticket-2"]}]},
+            }
+        },
+    )
+
+    # Test with experimental_use_trino=True
+    df = test_context_mock.foundry_sql_server_v2.query_foundry_sql(
+        "SELECT * FROM `ri.foundry.main.dataset.test-dataset`",
+        experimental_use_trino=True,
+    )
+
+    # Verify the query was modified to include the Trino backend hint
+    call_args = query_matcher.call_args_list[0]
+    request = call_args[0][0]
+    request_json = request.json()
+
+    assert "SELECT /*+ backend(trino) */ * FROM" in request_json["querySpec"]["query"]
+    assert df.shape[0] == 3
+
+    # Reset for second test
+    query_matcher.reset_mock()
+
+    # Test with experimental_use_trino=False (default)
+    df = test_context_mock.foundry_sql_server_v2.query_foundry_sql(
+        "SELECT * FROM `ri.foundry.main.dataset.test-dataset`",
+        experimental_use_trino=False,
+    )
+
+    # Verify the query was NOT modified
+    call_args = query_matcher.call_args_list[0]
+    request = call_args[0][0]
+    request_json = request.json()
+
+    assert request_json["querySpec"]["query"] == "SELECT * FROM `ri.foundry.main.dataset.test-dataset`"
+    assert "backend(trino)" not in request_json["querySpec"]["query"]
+    assert df.shape[0] == 3
+
+
+def test_v2_poll_for_query_completion_timeout(mocker, test_context_mock):
+    """Test that V2 query times out correctly when polling takes too long."""
+    mocker.patch("time.sleep")  # we do not want to wait in tests
+
+    # Mock the api_query endpoint (initial query execution)
+    test_context_mock.mock_adapter.register_uri(
+        "POST",
+        build_api_url(TEST_HOST.url, "foundry-sql-server", "sql-endpoint/v1/queries/query"),
+        json={"type": "running", "running": {"queryHandle": {"queryId": "test-query-timeout-123", "type": "foundry"}}},
+    )
+
+    # Mock the api_status endpoint to always return running status
+    test_context_mock.mock_adapter.register_uri(
+        "POST",
+        build_api_url(TEST_HOST.url, "foundry-sql-server", "sql-endpoint/v1/queries/status"),
+        json={"status": {"type": "running", "running": {}}},
+    )
+
+    with pytest.raises(FoundrySqlQueryClientTimedOutError):
+        test_context_mock.foundry_sql_server_v2.query_foundry_sql(
+            "SELECT * FROM `ri.foundry.main.dataset.test-dataset`",
+            timeout=0.001,
+        )
+
+
+def test_v2_ansi_dialect_not_supported(test_context_mock):
+    """Test that V2 client rejects ANSI SQL dialect."""
+    with pytest.raises(TypeError, match="'ANSI' is not a valid option for sql_dialect"):
+        test_context_mock.foundry_sql_server_v2.query_foundry_sql(
+            "SELECT * FROM `ri.foundry.main.dataset.test-dataset`",
+            sql_dialect="ANSI",  # type: ignore[arg-type]
+        )