From 2f19b0e8d865ab9c4d6ec27e0f4a11cd39b223fc Mon Sep 17 00:00:00 2001
From: timgl <tim@posthog.com>
Date: Fri, 25 Oct 2024 22:08:32 +0100
Subject: [PATCH] fix

---
 hogql_parser/README.md  |   2 +
 hogql_parser/parser.cpp | 108 +++++++++++++++++-----------------------
 2 files changed, 49 insertions(+), 61 deletions(-)

diff --git a/hogql_parser/README.md b/hogql_parser/README.md
index c8436358ffc9a..7ce505a412ce9 100644
--- a/hogql_parser/README.md
+++ b/hogql_parser/README.md
@@ -1,3 +1,5 @@
 # HogQL Parser
 
 Blazing fast HogQL parsing. This package can only work in the context of the PostHog Django app, as it imports from `posthog.hogql`.
+
+You can test changes locally by running `pip install ./hogql_parser`
\ No newline at end of file
diff --git a/hogql_parser/parser.cpp b/hogql_parser/parser.cpp
index 128f7006db93b..373905e43bc10 100644
--- a/hogql_parser/parser.cpp
+++ b/hogql_parser/parser.cpp
@@ -992,78 +992,64 @@ class HogQLParseTreeConverter : public HogQLParserBaseVisitor {
       }
     }
 
-    auto limit_and_offset_clause_ctx = ctx->limitAndOffsetClause();
-    if (limit_and_offset_clause_ctx) {
-      PyObject* limit;
+  auto limit_and_offset_clause_ctx = ctx->limitAndOffsetClause();
+  if (limit_and_offset_clause_ctx) {
+    PyObject* limit;
+    try {
+      limit = visitAsPyObject(limit_and_offset_clause_ctx->columnExpr(0));
+    } catch (...) {
+      Py_DECREF(select_query);
+      throw;
+    }
+    err_indicator = PyObject_SetAttrString(select_query, "limit", limit);
+    Py_DECREF(limit);
+    if (err_indicator == -1) {
+      Py_DECREF(select_query);
+      throw PyInternalError();
+    }
+    auto offset_ctx = limit_and_offset_clause_ctx->columnExpr(1);
+    if (offset_ctx) {
+      PyObject* offset;
       try {
-        limit = visitAsPyObject(limit_and_offset_clause_ctx->columnExpr(0));
+        offset = visitAsPyObject(offset_ctx);
       } catch (...) {
         Py_DECREF(select_query);
         throw;
       }
-      err_indicator = PyObject_SetAttrString(select_query, "limit", limit);
-      Py_DECREF(limit);
+      err_indicator = PyObject_SetAttrString(select_query, "offset", offset);
+      Py_DECREF(offset);
       if (err_indicator == -1) {
         Py_DECREF(select_query);
         throw PyInternalError();
       }
-      auto offset_ctx = limit_and_offset_clause_ctx->columnExpr(1);
-      if (offset_ctx) {
-        PyObject* offset;
-        try {
-          offset = visitAsPyObject(offset_ctx);
-        } catch (...) {
-          Py_DECREF(select_query);
-          throw;
-        }
-        err_indicator = PyObject_SetAttrString(select_query, "offset", offset);
-        Py_DECREF(offset);
-        if (err_indicator == -1) {
-          Py_DECREF(select_query);
-          throw PyInternalError();
-        }
-      }
-      auto limit_by_exprs_ctx = limit_and_offset_clause_ctx->columnExprList();
-      if (limit_by_exprs_ctx) {
-        PyObject* limit_by_exprs;
-        try {
-          limit_by_exprs = visitAsPyObject(limit_by_exprs_ctx);
-        } catch (...) {
-          Py_DECREF(select_query);
-          throw;
-        }
-        err_indicator = PyObject_SetAttrString(select_query, "limit_by", limit_by_exprs);
-        Py_DECREF(limit_by_exprs);
-        if (err_indicator == -1) {
-          Py_DECREF(select_query);
-          throw PyInternalError();
-        }
-      }
-      if (limit_and_offset_clause_ctx->WITH() && limit_and_offset_clause_ctx->TIES()) {
-        err_indicator = PyObject_SetAttrString(select_query, "limit_with_ties", Py_True);
-        if (err_indicator == -1) {
-          Py_DECREF(select_query);
-          throw PyInternalError();
-        }
-      }
-    } else {
-      auto offset_only_clause_ctx = ctx->offsetOnlyClause();
-      if (offset_only_clause_ctx) {
-        PyObject* offset_only_clause;
-        try {
-          offset_only_clause = visitAsPyObject(offset_only_clause_ctx->columnExpr());
-        } catch (...) {
-          Py_DECREF(select_query);
-          throw;
-        }
-        err_indicator = PyObject_SetAttrString(select_query, "offset", offset_only_clause);
-        Py_DECREF(offset_only_clause);
-        if (err_indicator == -1) {
-          Py_DECREF(select_query);
-          throw PyInternalError();
-        }
+    }
+    if (limit_and_offset_clause_ctx->WITH() && limit_and_offset_clause_ctx->TIES()) {
+      err_indicator = PyObject_SetAttrString(select_query, "limit_with_ties", Py_True);
+      if (err_indicator == -1) {
+        Py_DECREF(select_query);
+        throw PyInternalError();
       }
     }
+  }
+
+  // Handle limitByClause
+  auto limit_by_clause_ctx = ctx->limitByClause();
+  if (limit_by_clause_ctx) {
+    PyObject* limit_by_expr;
+    try {
+      limit_by_expr = visitAsPyObject(limit_by_clause_ctx);
+    } catch (...) {
+      Py_DECREF(select_query);
+      throw;
+    }
+    err_indicator = PyObject_SetAttrString(select_query, "limit_by", limit_by_expr);
+    Py_DECREF(limit_by_expr);
+    if (err_indicator == -1) {
+      Py_DECREF(select_query);
+      throw PyInternalError();
+    }
+  }
+
 
     auto array_join_clause_ctx = ctx->arrayJoinClause();
     if (array_join_clause_ctx) {