Support for LIMIT BY clause

Author: ods

aiochsa/__init__.py

@@ -3,6 +3,7 @@
 from .client import Client
 from .exc import DBException, ProtocolError
 from .pool import connect, create_pool
+from .sql import select
 
 try:
     __version__ = get_distribution(__name__).version

aiochsa/dialect.py

@@ -4,6 +4,8 @@
 from sqlalchemy.sql import crud
 from sqlalchemy.sql.compiler import SQLCompiler
 
+from .sql import _ORDER_BY_SENTINEL
+
 
 class ClickhouseSaSQLCompiler(ClickHouseCompiler):
 
@@ -12,6 +14,25 @@ class ClickhouseSaSQLCompiler(ClickHouseCompiler):
     def get_from_hint_text(self, table, text):
         return text
 
+    def order_by_clause(self, select, **kw):
+        text = ''
+        if select._order_by_clause is not _ORDER_BY_SENTINEL:
+            text = super().order_by_clause(select, **kw)
+
+        # Hack to add LIMIT BY clause
+        limit_by_clause = getattr(select, '_limit_by_clause', None)
+        if limit_by_clause is not None and limit_by_clause.clauses.clauses:
+            text += ' LIMIT '
+            if limit_by_clause.offset is not None:
+                text += f'{self.process(limit_by_clause.offset, **kw)}, '
+            text += self.process(limit_by_clause.limit, **kw)
+            limit_by_exprs = limit_by_clause.clauses._compiler_dispatch(
+                self, **kw,
+            )
+            text += f' BY {limit_by_exprs}'
+
+        return text
+
     def visit_insert(self, insert_stmt, asfrom=False, **kw):
         assert not self.stack # INSERT only at top level
 

aiochsa/sql.py

@@ -0,0 +1,63 @@
+from types import SimpleNamespace
+
+from clickhouse_sqlalchemy.sql import Select as _BaseSelect
+from sqlalchemy.sql.base import _generative
+from sqlalchemy.sql.elements import ClauseList, _literal_as_label_reference
+from sqlalchemy.sql.selectable import _offset_or_limit_clause
+
+
+class LimitByClause:
+
+    def __init__(self, clauses, offset, limit):
+        self.clauses = ClauseList(
+            *clauses, _literal_as_text=_literal_as_label_reference,
+        )
+        self.offset = _offset_or_limit_clause(offset)
+        self.limit = _offset_or_limit_clause(limit)
+
+
+
+# Minimal value to trick SQLAlchemy into believing that there is a ORDER BY
+# clause
+_ORDER_BY_SENTINEL = SimpleNamespace(clauses=True)
+
+
+class Select(_BaseSelect):
+
+    _limit_by_clause = None
+
+    @_generative
+    def limit_by(self, *clauses, offset=None, limit):
+        """return a new selectable with the given
+        `LIMIT [offset,] limit BY clauses` criterion applied.
+
+        e.g.::
+
+            stmt = (
+                select([table.c.name, table.c.amount])
+                    .order_by(table.c.amount.desc())
+                    .limit_by(table.c.name, limit=1)
+            )
+        """
+        # LIMIT BY clause comes just after ORDER BY, so we hack the later's
+        # machinery to add former too.  Alternative would be redefining
+        # whole `_compose_select_body` method just to insert 2 lines of code.
+        # The worst here is that this method is redefined in
+        # clickhouse_sqlalchemy and still is under active development.
+        if not self._order_by_clause.clauses:
+            self._order_by_clause = _ORDER_BY_SENTINEL
+
+        self._limit_by_clause = LimitByClause(clauses, offset, limit)
+
+    def append_order_by(self, *clauses):
+        if self._order_by_clause is _ORDER_BY_SENTINEL:
+            self._order_by_clause = ClauseList()
+        super().append_order_by(*clauses)
+        if (
+            self._limit_by_clause is not None and
+            not self._order_by_clause.clauses
+        ):
+            self._order_by_clause = _ORDER_BY_SENTINEL
+
+
+select = Select

tests/conftest.py

@@ -2,6 +2,7 @@
 from datetime import datetime
 from decimal import Decimal
 
+import clickhouse_sqlalchemy
 import pytest
 import sqlalchemy as sa
 
@@ -180,3 +181,12 @@ async def _create(sa_type):
         )
 
     return _create
+
+
+@pytest.fixture(params=[
+    sa.select,
+    clickhouse_sqlalchemy.select,
+    aiochsa.select,
+])
+def any_select(request):
+    return request.param