[iceberg] branch master updated: Spark 3.0: Ignore SQL comments checking for Iceberg commands (#4497)

[sz...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

szehon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iceberg.git


The following commit(s) were added to refs/heads/master by this push:
     new 0b48f51ee Spark 3.0: Ignore SQL comments checking for Iceberg commands (#4497)
0b48f51ee is described below

commit 0b48f51ee4aa2fb2c33b8d1335abb236a6a8d1c4
Author: Kyle Bendickson <kj...@gmail.com>
AuthorDate: Wed Apr 20 16:57:01 2022 -0700

    Spark 3.0: Ignore SQL comments checking for Iceberg commands (#4497)
---
 .../IcebergSparkSqlExtensionsParser.scala          | 10 +++++-
 .../spark/extensions/TestCallStatementParser.java  | 27 ++++++++++++++++
 .../extensions/TestExpireSnapshotsProcedure.java   | 36 ++++++++++++++++++++++
 3 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/spark/v3.0/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSparkSqlExtensionsParser.scala b/spark/v3.0/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSparkSqlExtensionsParser.scala
index b950bd704..6b6066773 100644
--- a/spark/v3.0/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSparkSqlExtensionsParser.scala
+++ b/spark/v3.0/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/IcebergSparkSqlExtensionsParser.scala
@@ -108,7 +108,15 @@ class IcebergSparkSqlExtensionsParser(delegate: ParserInterface) extends ParserI
   }
 
   private def isIcebergCommand(sqlText: String): Boolean = {
-    val normalized = sqlText.toLowerCase(Locale.ROOT).trim().replaceAll("\\s+", " ")
+    val normalized = sqlText.toLowerCase(Locale.ROOT).trim()
+        // Strip simple SQL comments that terminate a line, e.g. comments starting with `--`
+        .replaceAll("--.*?\\n", " ")
+        // Strip newlines.
+        .replaceAll("\\s+", " ")
+        // Strip comments of the form  /* ... */. This must come after stripping newlines so that
+        // comments that span multiple lines are caught.
+        .replaceAll("/\\*.*?\\*/", " ")
+        .trim()
     normalized.startsWith("call") || (
         normalized.startsWith("alter table") && (
             normalized.contains("add partition field") ||
diff --git a/spark/v3.0/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestCallStatementParser.java b/spark/v3.0/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestCallStatementParser.java
index 99e9d0621..c8c4316a1 100644
--- a/spark/v3.0/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestCallStatementParser.java
+++ b/spark/v3.0/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestCallStatementParser.java
@@ -22,8 +22,10 @@ package org.apache.iceberg.spark.extensions;
 import java.math.BigDecimal;
 import java.sql.Timestamp;
 import java.time.Instant;
+import java.util.List;
 import org.apache.iceberg.AssertHelpers;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.apache.spark.sql.SparkSession;
 import org.apache.spark.sql.catalyst.expressions.Expression;
 import org.apache.spark.sql.catalyst.expressions.Literal;
@@ -130,6 +132,31 @@ public class TestCallStatementParser {
     checkArg(call, 0, "value", DataTypes.StringType);
   }
 
+  @Test
+  public void testCallStripsComments() throws ParseException {
+    List<String> callStatementsWithComments = Lists.newArrayList(
+        "/* bracketed comment */  CALL cat.system.func('${spark.extra.prop}')",
+        "/**/  CALL cat.system.func('${spark.extra.prop}')",
+        "-- single line comment \n CALL cat.system.func('${spark.extra.prop}')",
+        "-- multiple \n-- single line \n-- comments \n CALL cat.system.func('${spark.extra.prop}')",
+        "/* select * from multiline_comment \n where x like '%sql%'; */ CALL cat.system.func('${spark.extra.prop}')",
+        "/* {\"app\": \"dbt\", \"dbt_version\": \"1.0.1\", \"profile_name\": \"profile1\", \"target_name\": \"dev\", " +
+            "\"node_id\": \"model.profile1.stg_users\"} \n*/ CALL cat.system.func('${spark.extra.prop}')",
+        "/* Some multi-line comment \n" +
+            "*/ CALL /* inline comment */ cat.system.func('${spark.extra.prop}') -- ending comment",
+        "CALL -- a line ending comment\n" +
+            "cat.system.func('${spark.extra.prop}')"
+    );
+    for (String sqlText : callStatementsWithComments) {
+      CallStatement call = (CallStatement) parser.parsePlan(sqlText);
+      Assert.assertEquals(ImmutableList.of("cat", "system", "func"), JavaConverters.seqAsJavaList(call.name()));
+
+      Assert.assertEquals(1, call.args().size());
+
+      checkArg(call, 0, "value", DataTypes.StringType);
+    }
+  }
+
   @Test
   public void testCallParseError() {
     AssertHelpers.assertThrows("Should fail with a sensible parse error", IcebergParseException.class,
diff --git a/spark/v3.0/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestExpireSnapshotsProcedure.java b/spark/v3.0/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestExpireSnapshotsProcedure.java
index 881ee5c1d..bc3e5e423 100644
--- a/spark/v3.0/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestExpireSnapshotsProcedure.java
+++ b/spark/v3.0/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestExpireSnapshotsProcedure.java
@@ -249,4 +249,40 @@ public class TestExpireSnapshotsProcedure extends SparkExtensionsTestBase {
         catalogName, currentTimestamp, tableIdent);
     assertEquals("Procedure output must match", ImmutableList.of(row(0L, 0L, 1L)), output);
   }
+
+  @Test
+  public void testExpireSnapshotsProcedureWorksWithSqlComments() {
+    // Ensure that systems such as dbt, that inject comments into the generated SQL files, will
+    // work with Iceberg-specific DDL
+    sql("CREATE TABLE %s (id bigint NOT NULL, data string) USING iceberg", tableName);
+
+    sql("INSERT INTO TABLE %s VALUES (1, 'a')", tableName);
+    sql("INSERT INTO TABLE %s VALUES (2, 'b')", tableName);
+
+    Table table = validationCatalog.loadTable(tableIdent);
+
+    Assert.assertEquals("Should be 2 snapshots", 2, Iterables.size(table.snapshots()));
+
+    waitUntilAfter(table.currentSnapshot().timestampMillis());
+
+    Timestamp currentTimestamp = Timestamp.from(Instant.ofEpochMilli(System.currentTimeMillis()));
+
+    String callStatement =
+        "/* CALL statement is used to expire snapshots */\n" +
+        "-- And we have single line comments as well \n" +
+        "/* And comments that span *multiple* \n" +
+        " lines */ CALL /* this is the actual CALL */ %s.system.expire_snapshots(" +
+        "   older_than => TIMESTAMP '%s'," +
+        "   table => '%s'," +
+        "   retain_last => 1)";
+    List<Object[]> output = sql(
+        callStatement, catalogName, currentTimestamp, tableIdent);
+    assertEquals("Procedure output must match",
+        ImmutableList.of(row(0L, 0L, 1L)),
+        output);
+
+    table.refresh();
+
+    Assert.assertEquals("Should be 1 snapshot remaining", 1, Iterables.size(table.snapshots()));
+  }
 }