Overwrite limit_clause in TrinoSQLCompiler

d-cole · Jan 14, 2022 · 0fb9d8f · 0fb9d8f
1 parent 62f53c9
commit 0fb9d8f
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 1 deletion.
diff --git a/tests/unit/sqlalchemy/test_compiler.py b/tests/unit/sqlalchemy/test_compiler.py
@@ -0,0 +1,46 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+from sqlalchemy import Table, MetaData, Column, Integer, String, select
+
+from trino.sqlalchemy.dialect import TrinoDialect
+
+metadata = MetaData()
+table = Table(
+    'table',
+    metadata,
+    Column('id', Integer, primary_key=True),
+    Column('name', String),
+)
+
+
+@pytest.fixture
+def dialect():
+    return TrinoDialect()
+
+
+def test_limit_offset(dialect):
+    statement = select(table).limit(10).offset(0)
+    query = statement.compile(dialect=dialect)
+    assert str(query) == 'SELECT "table".id, "table".name \nFROM "table"\nOFFSET :param_1\nLIMIT :param_2'
+
+
+def test_limit(dialect):
+    statement = select(table).limit(10)
+    query = statement.compile(dialect=dialect)
+    assert str(query) == 'SELECT "table".id, "table".name \nFROM "table"\nLIMIT :param_1'
+
+
+def test_offset(dialect):
+    statement = select(table).offset(0)
+    query = statement.compile(dialect=dialect)
+    assert str(query) == 'SELECT "table".id, "table".name \nFROM "table"\nOFFSET :param_1'
diff --git a/trino/sqlalchemy/compiler.py b/trino/sqlalchemy/compiler.py
@@ -91,7 +91,16 @@
 
 
 class TrinoSQLCompiler(compiler.SQLCompiler):
-    pass
+    def limit_clause(self, select, **kw):
+        """
+        Trino support only OFFSET...LIMIT but not LIMIT...OFFSET syntax.
+        """
+        text = ""
+        if select._offset_clause is not None:
+            text += "\nOFFSET " + self.process(select._offset_clause, **kw)
+        if select._limit_clause is not None:
+            text += "\nLIMIT " + self.process(select._limit_clause, **kw)
+        return text
 
 
 class TrinoDDLCompiler(compiler.DDLCompiler):