fix: BI-6846 Patch for TrinoSQLCompiler window functions (#1421)

* patch compile_ignore_nulls

* add test

* refactor test

* fix trino version

* move patch to package root

Author: khamitovdr

lib/dl_api_lib_testing/dl_api_lib_testing/connector/complex_queries.py

@@ -551,6 +551,7 @@ def test_window_functions(
                 "Date Sales": "SUM([Group Sales] WITHIN [order_date])",
                 "City Sales": "SUM([Group Sales] AMONG [order_date])",
                 "Total RSUM": 'RSUM([Group Sales], "asc" TOTAL)',
+                "First Round Count": "FIRST(ROUND(COUNT(1)))",
             },
         )
 
@@ -567,6 +568,7 @@ def test_window_functions(
                 ds.find_field(title="Date Sales"),
                 ds.find_field(title="City Sales"),
                 ds.find_field(title="Total RSUM"),
+                ds.find_field(title="First Round Count"),
             ],
             order_by=[
                 ds.find_field(title="order_date"),
@@ -585,7 +587,7 @@ def test_window_functions(
         assert {row[3] for row in data_rows}.issubset({str(i) for i in range(1, cnt + 1)})
 
         # There are as many [Unique Rank of Sales] values as there are rows
-        assert {row[4] for row in data_rows} == ({str(i) for i in range(1, cnt + 1)})
+        assert {row[4] for row in data_rows} == {str(i) for i in range(1, cnt + 1)}
 
         # [Rank of City Sales for Date] values are not greater than the number of [City] values
         assert len({row[5] for row in data_rows}) <= len({row[1] for row in data_rows})
@@ -603,6 +605,8 @@ def test_window_functions(
             # RSUM = previous RSUM value + value of current arg
             assert pytest.approx(float(data_rows[i][9])) == float(data_rows[i - 1][9]) + float(data_rows[i][2])
 
+        assert all(float(row[10]) == 1 for row in data_rows)
+
 
 class DefaultBasicNativeFunctionTestSuite(
     RegulatedTestCase, DataApiTestBase, DatasetTestBase, DbServiceFixtureTextClass

lib/dl_connector_trino/dl_connector_trino/__init__.py

@@ -0,0 +1 @@
+from dl_connector_trino import vendor_patches  # noqa: F401

lib/dl_connector_trino/dl_connector_trino/vendor_patches.py

@@ -0,0 +1,23 @@
+from typing import Any
+
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.compiler import SQLCompiler
+from trino.sqlalchemy.compiler import TrinoSQLCompiler
+
+
+# This is a temporary patch to fix https://github.com/trinodb/trino-python-client/pull/586
+# BI-6846
+@compiles(TrinoSQLCompiler.FirstValue)
+@compiles(TrinoSQLCompiler.LastValue)
+@compiles(TrinoSQLCompiler.NthValue)
+@compiles(TrinoSQLCompiler.Lead)
+@compiles(TrinoSQLCompiler.Lag)
+def compile_ignore_nulls(
+    element: TrinoSQLCompiler.GenericIgnoreNulls,
+    compiler: SQLCompiler,
+    **kwargs: Any,
+) -> str:
+    compiled = f"{element.name}({compiler.process(element.clauses, **kwargs)})"
+    if element.ignore_nulls:
+        compiled += " IGNORE NULLS"
+    return compiled

metapkg/poetry.lock

@@ -83,6 +83,7 @@ sqlalchemy = "==1.4.46, <2.0"
 sqlalchemy-bigquery = "==1.9.0"
 tabulate = "==0.9.0"
 tornado = "==6.4.2"
+trino = {extras = ["sqlalchemy"], version = "==0.331.0"}
 typeguard = "==4.1.5"
 typing-extensions = "==4.15.0"
 ujson = "==1.35"