Merge pull request #16073 from cdapio/spanner-query-fix

sidhdirenge · web-flow · commit 8e6f5a43db1b · 2025-12-31T12:36:45.000+05:30
[CDAP-21220] Fix range Where clause in the SpannerStructuredTable queries
diff --git a/cdap-storage-ext-spanner/pom.xml b/cdap-storage-ext-spanner/pom.xml
@@ -82,6 +82,11 @@
       <artifactId>snappy-java</artifactId>
       <version>1.1.10.7</version>
     </dependency>
+    <dependency>
+      <groupId>pl.pragmatists</groupId>
+      <artifactId>JUnitParams</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <profiles>
diff --git a/cdap-storage-ext-spanner/src/main/java/io/cdap/cdap/storage/spanner/SpannerStructuredTable.java b/cdap-storage-ext-spanner/src/main/java/io/cdap/cdap/storage/spanner/SpannerStructuredTable.java
@@ -29,6 +29,7 @@
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.TransactionContext;
 import com.google.cloud.spanner.Value;
+import com.google.common.annotations.VisibleForTesting;
 import io.cdap.cdap.api.dataset.lib.AbstractCloseableIterator;
 import io.cdap.cdap.api.dataset.lib.CloseableIterator;
 import io.cdap.cdap.spi.data.FieldSizeLimitExceededException;
@@ -545,39 +546,82 @@ private String getRangeWhereClause(Range range, Map<String, Value> parameters) {
     List<String> conditions = new ArrayList<>();
     int paramIndex = parameters.size();
 
-    int beginIndex = 1;
-    int beginFieldsCount = range.getBegin().size();
-    for (Field<?> field : range.getBegin()) {
-      // Edge case for = when we filter by the last field
-      // This is for case like ([KEY: "ns1", KEY2: 123], EXCLUSIVE, [KEY: "ns1", KEY2: 250], INCLUSIVE)
-      // We need to check the ending bound for KEY2, not KEY
-      String symbol =
-          beginIndex == beginFieldsCount && range.getBeginBound().equals(Range.Bound.EXCLUSIVE)
-              ? " > " : " >= ";
-      conditions.add(escapeName(field.getName()) + symbol + "@p_" + paramIndex);
-      parameters.put("p_" + paramIndex, getValue(field));
-      paramIndex++;
-      beginIndex++;
-    }
-
-    int endIndex = 1;
-    int endFieldsCount = range.getEnd().size();
-    for (Field<?> field : range.getEnd()) {
-      // Edge case for = when we filter by the last field
-      // This is for case like ([KEY: "ns1", KEY2: 123], INCLUSIVE, [KEY: "ns1", KEY2: 250], EXCLUSIVE)
-      // We need to check the ending bound for KEY2, not KEY
-      String symbol =
-          endIndex == endFieldsCount && range.getEndBound().equals(Range.Bound.EXCLUSIVE) ? " < "
-              : " <= ";
-      conditions.add(escapeName(field.getName()) + symbol + "@p_" + paramIndex);
-      parameters.put("p_" + paramIndex, getValue(field));
-      paramIndex++;
-      endIndex++;
+    if (!range.getBegin().isEmpty()) {
+      conditions.add(getCompositeKeyCondition(new ArrayList<>(range.getBegin()), range.getBeginBound(),
+                                                      true, parameters, paramIndex));
+      paramIndex += range.getBegin().size();
+    }
+
+    if (!range.getEnd().isEmpty()) {
+      conditions.add(getCompositeKeyCondition(new ArrayList<>(range.getEnd()), range.getEndBound(),
+                                                      false, parameters, paramIndex));
     }
 
     return String.join(" AND ", conditions);
   }
 
+  /**
+   * Builds a SQL fragment for one bound (lower or upper) of a range on a compound key,
+   * handling lexicographical ordering.
+   *
+   * <p>Example: For a lower bound `(Key1, Key2) >= ('A', 10)`, this generates:
+   * `((`Key1` > @p_0) OR (`Key1` = @p_0 AND `Key2` >= @p_1))`.</p>
+   * The parameters map will contain mappings for `p_0` to 'A' and `p_1` to 10.
+   *
+   * @param fields The fields forming the compound key (e.g., namespace, application, version).
+   * @param bound The bound type (INCLUSIVE or EXCLUSIVE).
+   * @param isLowerBound True if building the condition for the beginning of the range (lower bound),
+   *                     false if for the end of the range (upper bound).
+   * @param parameters The map to add Spanner parameters to. Parameter names will be generated
+   *                   starting from the current size of this map.
+   * @return A SQL fragment for the composite range bound.
+   */
+  @VisibleForTesting
+  String getCompositeKeyCondition(List<Field<?>> fields, Range.Bound bound, boolean isLowerBound,
+                                      Map<String, Value> parameters, int startParamIndex) {
+    if (fields.isEmpty()) {
+      return "";
+    }
+
+    List<String> orClauses = new ArrayList<>();
+    StringBuilder equalityPrefix = new StringBuilder();
+
+    for (int i = 0; i < fields.size(); i++) {
+      Field<?> field = fields.get(i);
+      String paramName = "p_" + (startParamIndex + i);
+      parameters.put(paramName, getValue(field));
+
+      // 1. Determine operator: strictly >/< for non-terminal fields,
+      // or based on inclusivity for the final field.
+      boolean isLast = (i == fields.size() - 1);
+      boolean isInclusive = bound == Range.Bound.INCLUSIVE;
+      String op;
+      if (isLowerBound) {
+        op = (isLast && isInclusive) ? " >= " : " > ";
+      } else {
+        op = (isLast && isInclusive) ? " <= " : " < ";
+      }
+
+      // 2. Construct the OR clause for this depth
+      String escapedName = escapeName(field.getName());
+      String condition = escapedName + op + "@" + paramName;
+      if (equalityPrefix.length() > 0) {
+        orClauses.add("(" + equalityPrefix + " AND " + condition + ")");
+      } else {
+        orClauses.add("(" + condition + ")");
+      }
+
+      if (!isLast) {
+        if (equalityPrefix.length() > 0) {
+          equalityPrefix.append(" AND ");
+        }
+        equalityPrefix.append(escapedName).append(" = @").append(paramName);
+      }
+    }
+
+    return "(" + String.join(" OR ", orClauses) + ")";
+  }
+
   private String getFieldsWhereClause(Collection<Field<?>> fields, Map<String, Value> parameters) {
     List<String> conditions = new ArrayList<>();
     int paramIndex = parameters.size();
diff --git a/cdap-storage-ext-spanner/src/test/java/io/cdap/cdap/storage/spanner/SpannerStructuredTableTest.java b/cdap-storage-ext-spanner/src/test/java/io/cdap/cdap/storage/spanner/SpannerStructuredTableTest.java
@@ -16,20 +16,32 @@
 
 package io.cdap.cdap.storage.spanner;
 
+import com.google.cloud.spanner.Value;
 import io.cdap.cdap.api.metrics.MetricsCollector;
 import io.cdap.cdap.spi.data.StorageProviderContext;
 import io.cdap.cdap.spi.data.StructuredTable;
 import io.cdap.cdap.spi.data.StructuredTableAdmin;
 import io.cdap.cdap.spi.data.StructuredTableTest;
+import io.cdap.cdap.spi.data.table.field.Field;
+import io.cdap.cdap.spi.data.table.field.Fields;
+import io.cdap.cdap.spi.data.table.field.Range;
 import io.cdap.cdap.spi.data.transaction.TransactionRunner;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import junitparams.JUnitParamsRunner;
+import junitparams.Parameters;
+import junitparams.naming.TestCaseName;
 import org.junit.AfterClass;
+import org.junit.Assert;
 import org.junit.Assume;
 import org.junit.BeforeClass;
 import org.junit.Ignore;
+import org.junit.Test;
+import org.junit.runner.RunWith;
 
 /**
  * Unit tests for GCP spanner implementation of the {@link StructuredTable}. This test needs the following
@@ -43,6 +55,7 @@
  *   json that has the "Cloud Spanner Database User" role</li>
  * </ul>
  */
+@RunWith(JUnitParamsRunner.class)
 public class SpannerStructuredTableTest extends StructuredTableTest {
 
   private static SpannerStorageProvider storageProvider;
@@ -93,6 +106,113 @@ protected TransactionRunner getTransactionRunner() {
     return storageProvider.getTransactionRunner();
   }
 
+  @Test
+  @Parameters(method = "getCompositeKeyConditionScenarios")
+  @TestCaseName("{0}")
+  public void testGetCompositeKeyCondition(String description, List<Field<?>> fields,
+      Range.Bound bound, boolean isLowerBound, int startParamIndex, String expectedSql,
+      Map<String, Value> expectedParams) {
+    SpannerStructuredTableSchema schema = new SpannerStructuredTableSchema(SIMPLE_SPEC.getTableId(),
+        SIMPLE_SPEC.getFieldTypes(), SIMPLE_SPEC.getPrimaryKeys(), SIMPLE_SPEC.getIndexes(), null);
+    try (SpannerStructuredTable table = new SpannerStructuredTable(null, schema, null)) {
+      Map<String, Value> actualParams = new HashMap<>();
+      String actualSql = table.getCompositeKeyCondition(fields, bound, isLowerBound, actualParams,
+          startParamIndex);
+      Assert.assertEquals("SQL mismatch in: " + description, expectedSql, actualSql);
+      Assert.assertEquals("Params mismatch in: " + description, expectedParams, actualParams);
+    }
+  }
+
+  private Object[] getCompositeKeyConditionScenarios() {
+    return new Object[]{
+        // Format: Name | Fields | Bound Type | isLower | startParamIndex | Expected SQL | Expected Params
+        new Object[]{
+            "COMPOSITE: (col1, col2) >= ('A', 5) [Lower Inclusive]",
+            Arrays.asList(Fields.stringField("col1", "A"), Fields.intField("col2", 5)),
+            Range.Bound.INCLUSIVE, true, 0,
+            "((`col1` > @p_0) OR (`col1` = @p_0 AND `col2` >= @p_1))",
+            params("p_0", Value.string("A"), "p_1", Value.int64(5))
+        },
+        new Object[]{
+            "COMPOSITE: (col1, col2) > ('A', 5) [Lower Exclusive]",
+            Arrays.asList(Fields.stringField("col1", "A"), Fields.intField("col2", 5)),
+            Range.Bound.EXCLUSIVE, true, 0,
+            "((`col1` > @p_0) OR (`col1` = @p_0 AND `col2` > @p_1))",
+            params("p_0", Value.string("A"), "p_1", Value.int64(5))
+        },
+        new Object[]{
+            "COMPOSITE: (col1, col2) <= ('B', 10) [Upper Inclusive]",
+            Arrays.asList(Fields.stringField("col1", "B"), Fields.intField("col2", 10)),
+            Range.Bound.INCLUSIVE, false, 0,
+            "((`col1` < @p_0) OR (`col1` = @p_0 AND `col2` <= @p_1))",
+            params("p_0", Value.string("B"), "p_1", Value.int64(10))
+        },
+
+        new Object[]{
+            "SINGLE: col1 >= 'A' [Lower Inclusive]",
+            Collections.singletonList(Fields.stringField("col1", "A")),
+            Range.Bound.INCLUSIVE, true, 0,
+            "((`col1` >= @p_0))",
+            params("p_0", Value.string("A"))
+        },
+        new Object[]{
+            "SINGLE: col1 < 'Z' [Upper Exclusive]",
+            Collections.singletonList(Fields.stringField("col1", "Z")),
+            Range.Bound.EXCLUSIVE, false, 0,
+            "((`col1` < @p_0))",
+            params("p_0", Value.string("Z"))
+        },
+
+        new Object[]{
+            "DEEP: (k1, k2, k3, k4) <= ('x', 1, true, 99) [Upper Inclusive]",
+            Arrays.asList(
+                Fields.stringField("k1", "x"),
+                Fields.intField("k2", 1),
+                Fields.booleanField("k3", true),
+                Fields.longField("k4", 99L)
+            ),
+            Range.Bound.INCLUSIVE, false, 0,
+            "((`k1` < @p_0) OR " +
+                "(`k1` = @p_0 AND `k2` < @p_1) OR " +
+                "(`k1` = @p_0 AND `k2` = @p_1 AND `k3` < @p_2) OR " +
+                "(`k1` = @p_0 AND `k2` = @p_1 AND `k3` = @p_2 AND `k4` <= @p_3))",
+            params("p_0", Value.string("x"), "p_1", Value.int64(1), "p_2", Value.bool(true), "p_3",
+                Value.int64(99L))
+        },
+
+        new Object[]{
+            "EDGE: Empty field list returns empty string",
+            Collections.emptyList(),
+            Range.Bound.INCLUSIVE, true, 0,
+            "",
+            Collections.emptyMap()
+        },
+
+        new Object[]{
+            "OFFSET: 1 Field, Lower Inclusive, startParamIndex = 5",
+            Collections.singletonList(Fields.stringField("id", "123")),
+            Range.Bound.INCLUSIVE, true, 5,
+            "((`id` >= @p_5))",
+            params("p_5", Value.string("123"))
+        },
+        new Object[]{
+            "OFFSET: 2 Fields, Upper Exclusive, startParamIndex = 10",
+            Arrays.asList(Fields.stringField("col1", "X"), Fields.intField("col2", 99)),
+            Range.Bound.EXCLUSIVE, false, 10,
+            "((`col1` < @p_10) OR (`col1` = @p_10 AND `col2` < @p_11))",
+            params("p_10", Value.string("X"), "p_11", Value.int64(99))
+        }
+    };
+  }
+
+  private static Map<String, Value> params(Object... entries) {
+    Map<String, Value> map = new HashMap<>();
+    for (int i = 0; i < entries.length; i += 2) {
+      map.put((String) entries[i], (Value) entries[i + 1]);
+    }
+    return map;
+  }
+
   private static final class MockStorageProviderContext implements StorageProviderContext {
 
     private final Map<String, String> config;
