Implemented LIMIT clause support for Athena

Author: bilaliqbalr

src/Bilaliqbalr/Athena/src/Connection.php

@@ -10,6 +10,7 @@
 use Illuminate\Database\QueryException;
 use Bilaliqbalr\Athena\Query\Grammar as QueryGrammar;
 use Bilaliqbalr\Athena\Query\Processor;
+use Illuminate\Support\Arr;
 use Illuminate\Support\Facades\Config;
 use Bilaliqbalr\Athena\Schema\Builder;
 use Bilaliqbalr\Athena\Schema\Grammar as SchemaGrammar;
@@ -181,11 +182,11 @@ public function formatCSVFileQueryResults($filePath)
     }
 
     /**
-     * @param \Illuminate\Database\Query\Builder $builder
      * @param null $query
      * @param null $binding
      *
      * @return mixed
+     * @throws Exception
      */
     protected function prepareQuery($query, $binding)
     {
@@ -197,6 +198,29 @@ protected function prepareQuery($query, $binding)
             }
         }
 
+        // Modifying query & preparing it for LIMIT as per Athena
+        if (is_int(stripos($query, 'BETWEENLIMIT'))) {
+            // Checking if ROW_NUMBER() OVER() window function applied, then take it as LIMIT query
+            if (!is_int(stripos($query, 'ROW_NUMBER()')) || !is_int(stripos($query, ' rn '))) {
+                throw new Exception("Error: Required `ROW_NUMBER() OVER(...) as rn` to implement LIMIT functionality");
+            } else {
+                $queryParts = preg_split("/BETWEENLIMIT/i", $query);
+                preg_match_all('!\d+!', array_pop($queryParts), $matches);
+                // Calculating offset and limit for Athena
+                $perPage = ($matches[0][0] - 1);
+
+                // Only apply this limit if we have per page greater than 0.
+                // This prevent BETWEEN as WHERE clause to be treated like LIMIT & OFFSET, which occurs if we have
+                // both BETWEEN and ROW_NUMBER() in query but that no LIMIT
+                if ($perPage > 0) {
+                    $page = ($matches[0][1] / $perPage) + 1;
+                    $from = ($perPage * ($page - 1)) + 1;
+                    $to = ($perPage * $page);
+                    $query = "SELECT * FROM ( " . Arr::first($queryParts) . " ) WHERE rn BETWEEN $from AND $to";
+                }
+            }
+        }
+
         return str_replace('`', '', $query);
     }
 
@@ -216,7 +240,7 @@ public function getDownloadedFilePath()
      * @return array|\Aws\Result
      * @throws Exception
      */
-    protected function executeQuery($query, $bindings)
+    protected function executeQuery(&$query, $bindings)
     {
         $query = $this->prepareQuery($query, $bindings);
 
@@ -275,7 +299,7 @@ public function statement($query, $bindings = [])
         $this->executeQuery($query, $bindings);
 
         $this->logQuery(
-            $query, $bindings, $this->getElapsedTime($start)
+            $query, [], $this->getElapsedTime($start)
         );
 
         return true;
@@ -306,12 +330,11 @@ public function select($query, $bindings = [], $useReadPdo = true)
 
             if ($this->downloadFileFromS3ToLocalServer($s3FilePath, $localFilePath)) {
                 $this->localFilePath = $localFilePath;
-                dd($s3FilePath, $this->localFilePath);
                 $result = $this->formatCSVFileQueryResults($this->localFilePath);
             }
         }
         $this->logQuery(
-            $query, $bindings, $this->getElapsedTime($start)
+            $query, [], $this->getElapsedTime($start)
         );
 
         return $result;

src/Bilaliqbalr/Athena/src/Query/Grammar.php

@@ -15,7 +15,14 @@ class Grammar extends MySqlGrammar
      */
     protected function compileLimit(Builder $query, $limit)
     {
-        return 'BETWEEN '.(int) $limit;
+        // Only apply BETWEEN clause, if missing OFFSET, otherwise use presto way to LIMIT records
+        if (is_int($query->offset)) {
+            // using custom BETWEENLIMIT clause only to detect if it is limit to prevent conflict with BETWEEN.
+            // Handling it in Connection.php
+            return 'BETWEENLIMIT '.(int) $limit;
+        } else {
+            return parent::compileLimit($query, $limit);
+        }
     }
 
     /**