From d6c9779984057d46fb747d257447b8173c26bf8d Mon Sep 17 00:00:00 2001
From: Amit Dutta <adutta@fb.com>
Date: Thu, 17 Mar 2022 21:28:45 -0700
Subject: [PATCH] Throw PrestoException from VisitPlan.

---
 .../iterative/rule/ImplementOffset.java       | 13 +++++-----
 .../iterative/rule/TestImplementOffset.java   | 24 +++++++++++++++++++
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ImplementOffset.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ImplementOffset.java
index caa9880d3a621..3637ec3c66a27 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ImplementOffset.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ImplementOffset.java
@@ -13,9 +13,9 @@
  */
 package com.facebook.presto.sql.planner.iterative.rule;
 
-import com.facebook.presto.Session;
 import com.facebook.presto.matching.Captures;
 import com.facebook.presto.matching.Pattern;
+import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.plan.FilterNode;
 import com.facebook.presto.spi.plan.ProjectNode;
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
@@ -30,6 +30,7 @@
 
 import static com.facebook.presto.SystemSessionProperties.isOffsetClauseEnabled;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
 import static com.facebook.presto.sql.analyzer.ExpressionTreeUtils.createSymbolReference;
 import static com.facebook.presto.sql.planner.plan.AssignmentUtils.identityAssignmentsAsSymbolReferences;
 import static com.facebook.presto.sql.planner.plan.Patterns.offset;
@@ -63,15 +64,13 @@ public Pattern<OffsetNode> getPattern()
         return PATTERN;
     }
 
-    @Override
-    public boolean isEnabled(Session session)
-    {
-        return isOffsetClauseEnabled(session);
-    }
-
     @Override
     public Result apply(OffsetNode parent, Captures captures, Context context)
     {
+        if (!isOffsetClauseEnabled(context.getSession())) {
+            throw new PrestoException(NOT_SUPPORTED, "Offset support is not enabled");
+        }
+
         VariableReferenceExpression rowNumberSymbol = context.getVariableAllocator().newVariable("row_number", BIGINT);
 
         RowNumberNode rowNumberNode = new RowNumberNode(
diff --git a/presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestImplementOffset.java b/presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestImplementOffset.java
index daac159efb9ed..9fb7b577f732c 100644
--- a/presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestImplementOffset.java
+++ b/presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestImplementOffset.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.sql.planner.iterative.rule;
 
+import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
 import com.facebook.presto.sql.planner.assertions.ExpressionMatcher;
 import com.facebook.presto.sql.planner.assertions.RowNumberSymbolMatcher;
@@ -57,6 +58,29 @@ public void testReplaceOffsetOverValues()
                                                 .withAlias("row_num", new RowNumberSymbolMatcher()))));
     }
 
+    @Test(expectedExceptions = PrestoException.class, expectedExceptionsMessageRegExp = "Offset support is not enabled")
+    public void testOffsetClauseDisabled()
+    {
+        tester().assertThat(new ImplementOffset())
+                .on(p -> {
+                    VariableReferenceExpression a = p.variable("a");
+                    VariableReferenceExpression b = p.variable("b");
+                    return p.offset(
+                            2,
+                            p.values(a, b));
+                })
+                .matches(
+                        strictProject(
+                                ImmutableMap.of("a", new ExpressionMatcher("a"), "b", new ExpressionMatcher("b")),
+                                filter(
+                                        "row_num > BIGINT '2'",
+                                        rowNumber(
+                                                pattern -> pattern
+                                                        .partitionBy(ImmutableList.of()),
+                                                values("a", "b"))
+                                                .withAlias("row_num", new RowNumberSymbolMatcher()))));
+    }
+
     @Test
     public void testReplaceOffsetOverSort()
     {