Add limit offset caching support

Author: jwaterworth

sqlalchemy_redshift/dialect.py

@@ -4,7 +4,7 @@
 from packaging.version import Version
 import pkg_resources
 import sqlalchemy as sa
-from sqlalchemy import inspect
+from sqlalchemy import inspect, exc as sa_exc
 from sqlalchemy.dialects.postgresql.base import (
     PGCompiler, PGDDLCompiler, PGIdentifierPreparer, PGTypeCompiler,
     PGExecutionContext, PGDialect
@@ -560,7 +560,7 @@ class RedshiftDialectMixin(DefaultDialect):
     max_identifier_length = 127
     # explicitly disables statement cache to disable warnings in logs
     # ref: https://docs.sqlalchemy.org/en/14/core/connections.html#caching-for-third-party-dialects  # noqa
-    supports_statement_cache = False
+    supports_statement_cache = True
 
     statement_compiler = RedshiftCompiler
     ddl_compiler = RedshiftDDLCompiler
@@ -1053,16 +1053,49 @@ class RedshiftDialect_redshift_connector(RedshiftDialectMixin, PGDialect):
 
     class RedshiftCompiler_redshift_connector(RedshiftCompiler, PGCompiler):
         def limit_clause(self, select, **kw):
-            text = ""
-            if select._limit_clause is not None:
-                # an integer value for limit is retrieved
-                text += " \n LIMIT " + str(select._limit)
-            if select._offset_clause is not None:
-                if select._limit_clause is None:
-                    text += "\n LIMIT ALL"
-                # an integer value for offset is retrieved
-                text += " OFFSET " + str(select._offset)
-            return text
+            if sa_version >= Version('1.4.0'):
+                text = ""
+
+                limit_clause = select._limit_clause
+                offset_clause = select._offset_clause
+
+                if select._simple_int_clause(limit_clause):
+                    text += " \n LIMIT %s" % (
+                        self.process(
+                            limit_clause.render_literal_execute(),
+                            **kw
+                        )
+                    )
+                elif limit_clause is not None:
+                    raise sa_exc.CompileError(
+                        "dialect 'redshift-dialect' can only \
+                             render simple integers for LIMIT"
+                    )
+                if select._simple_int_clause(offset_clause):
+                    text += " \n OFFSET %s" % (
+                        self.process(
+                            offset_clause.render_literal_execute(),
+                            **kw
+                        )
+                    )
+                elif offset_clause is not None:
+                    raise sa_exc.CompileError(
+                        "dialect 'redshift-dialect' can only \
+                            render simple integers for OFFSET"
+                    )
+
+                return text
+            else:
+                text = ""
+                if select._limit_clause is not None:
+                    # an integer value for limit is retrieved
+                    text += " \n LIMIT " + str(select._limit)
+                if select._offset_clause is not None:
+                    if select._limit_clause is None:
+                        text += "\n LIMIT ALL"
+                    # an integer value for offset is retrieved
+                    text += " OFFSET " + str(select._offset)
+                return text
 
         def visit_mod_binary(self, binary, operator, **kw):
             return (

tests/test_limit_offset.py

@@ -0,0 +1,55 @@
+__author__ = "James Waterworth"
+"""
+Tests to valdiate that LIMIT and OFFSET values are cached
+at the appropriate stage of query generation
+"""
+
+import sqlalchemy as sa
+from packaging.version import Version
+
+from rs_sqla_test_utils.utils import clean, compile_query
+
+sa_version = Version(sa.__version__)
+
+meta = sa.MetaData()
+
+items = sa.Table(
+    "items",
+    meta,
+    sa.Column("id", sa.Integer, primary_key=True, autoincrement=False),
+    sa.Column("order_id", sa.Integer),
+    sa.Column("product_id", sa.Integer),
+    sa.Column("name", sa.String(255)),
+    sa.Column("qty", sa.Numeric(12, 4)),
+    sa.Column("price", sa.Numeric(12, 4)),
+    sa.Column("total_invoiced", sa.Numeric(12, 4)),
+    sa.Column("discount_invoiced", sa.Numeric(12, 4)),
+    sa.Column("grandtotal_invoiced", sa.Numeric(12, 4)),
+    sa.Column("created_at", sa.DateTime),
+    sa.Column("updated_at", sa.DateTime),
+)
+
+
+def test_select_1_item(stub_redshift_dialect):
+    query = sa.select(items.c.id).select_from(items).limit(1).offset(0)
+
+    assert (
+        clean(compile_query(query, stub_redshift_dialect))
+        == "SELECT items.id FROM items LIMIT 1 OFFSET 0"
+    )
+
+
+def test_limit_offset_values_not_cached(stub_redshift_dialect):
+    query = sa.select(items.c.id).select_from(items).limit(1).offset(0)
+
+    # compile our first query to get it cached
+    compile_query(query, stub_redshift_dialect)
+
+    second_query = sa.select(items.c.id).select_from(items).limit(2).offset(1)
+
+    # perform the same query, it should use the cached version but add the
+    # limit, offset values seperately
+    assert (
+        clean(compile_query(second_query, stub_redshift_dialect))
+        == "SELECT items.id FROM items LIMIT 2 OFFSET 1"
+    )