Merge branch 'apache:master' into master

Author: rakeshcusat

CHANGES.md

@@ -58,6 +58,7 @@
 
 ## Highlights
 
+* Managed API for [Java](https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/managed/Managed.html) and [Python](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.managed.html#module-apache_beam.transforms.managed) supports [key I/O connectors](https://beam.apache.org/documentation/io/connectors/) Iceberg, Kafka, and BigQuery.
 * New highly anticipated feature X added to Python SDK ([#X](https://github.com/apache/beam/issues/X)).
 * New highly anticipated feature Y added to Java SDK ([#Y](https://github.com/apache/beam/issues/Y)).
 
@@ -71,6 +72,7 @@
 * Support custom coders in Reshuffle ([#29908](https://github.com/apache/beam/issues/29908), [#33356](https://github.com/apache/beam/issues/33356)).
 * [Java] Upgrade SLF4J to 2.0.16. Update default Spark version to 3.5.0. ([#33574](https://github.com/apache/beam/pull/33574))
 * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* Managed API for [Java](https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/managed/Managed.html) and [Python](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.managed.html#module-apache_beam.transforms.managed) supports [key I/O connectors](https://beam.apache.org/documentation/io/connectors/) Iceberg, Kafka, and BigQuery.
 
 ## Breaking Changes
 

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Lineage.java

@@ -17,11 +17,16 @@
  */
 package org.apache.beam.sdk.metrics;
 
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Set;
 import java.util.regex.Pattern;
 import org.apache.beam.sdk.annotations.Internal;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
@@ -32,12 +37,13 @@ public class Lineage {
   public static final String LINEAGE_NAMESPACE = "lineage";
   private static final Lineage SOURCES = new Lineage(Type.SOURCE);
   private static final Lineage SINKS = new Lineage(Type.SINK);
-  private static final Pattern RESERVED_CHARS = Pattern.compile("[:\\s.]");
+  // Reserved characters are backtick, colon, whitespace (space, \t, \n) and dot.
+  private static final Pattern RESERVED_CHARS = Pattern.compile("[:\\s.`]");
 
-  private final StringSet metric;
+  private final BoundedTrie metric;
 
   private Lineage(Type type) {
-    this.metric = Metrics.stringSet(LINEAGE_NAMESPACE, type.toString());
+    this.metric = Metrics.boundedTrie(LINEAGE_NAMESPACE, type.toString());
   }
 
   /** {@link Lineage} representing sources and optionally side inputs. */
@@ -50,111 +56,135 @@ public static Lineage getSinks() {
     return SINKS;
   }
 
-  /**
-   * Wrap segment to valid segment name.
-   *
-   * <p>Specifically, If there are reserved chars (colon, whitespace, dot), escape with backtick. If
-   * the segment is already wrapped, return the original.
-   *
-   * <p>This helper method is for internal and testing usage only.
-   */
-  @Internal
-  public static String wrapSegment(String value) {
-    if (value.startsWith("`") && value.endsWith("`")) {
-      return value;
+  @VisibleForTesting
+  static Iterable<String> getFQNParts(
+      String system,
+      @Nullable String subtype,
+      Iterable<String> segments,
+      @Nullable String lastSegmentSep) {
+
+    List<String> parts = new ArrayList<>();
+    parts.add(system + ":");
+    if (subtype != null) {
+      parts.add(subtype + ":");
     }
-    if (RESERVED_CHARS.matcher(value).find()) {
-      return String.format("`%s`", value);
-    }
-    return value;
-  }
 
-  /**
-   * Assemble fully qualified name (<a
-   * href="https://cloud.google.com/data-catalog/docs/fully-qualified-names">FQN</a>). Format:
-   *
-   * <ul>
-   *   <li>{@code system:segment1.segment2}
-   *   <li>{@code system:subtype:segment1.segment2}
-   *   <li>{@code system:`segment1.with.dots:clons`.segment2}
-   * </ul>
-   *
-   * <p>This helper method is for internal and testing usage only.
-   */
-  @Internal
-  public static String getFqName(
-      String system, @Nullable String subtype, Iterable<String> segments) {
-    StringBuilder builder = new StringBuilder(system);
-    if (!Strings.isNullOrEmpty(subtype)) {
-      builder.append(":").append(subtype);
-    }
-    int idx = 0;
-    for (String segment : segments) {
-      if (idx == 0) {
-        builder.append(":");
-      } else {
-        builder.append(".");
+    if (segments != null) {
+      Iterator<String> iterator = segments.iterator();
+      String previousSegment = null;
+      while (iterator.hasNext()) {
+        if (previousSegment != null) {
+          parts.add(wrapSegment(previousSegment) + ".");
+        }
+        previousSegment = iterator.next();
+      }
+
+      if (previousSegment != null) {
+        if (lastSegmentSep != null) {
+          List<String> subSegments =
+              Splitter.onPattern(lastSegmentSep).splitToList(wrapSegment(previousSegment));
+          for (int i = 0; i < subSegments.size() - 1; i++) {
+            parts.add(subSegments.get(i) + lastSegmentSep);
+          }
+          parts.add(subSegments.get(subSegments.size() - 1));
+        } else {
+          parts.add(wrapSegment(previousSegment));
+        }
       }
-      builder.append(wrapSegment(segment));
-      ++idx;
     }
-    return builder.toString();
+
+    return parts;
   }
 
   /**
-   * Assemble the FQN of given system, and segments.
-   *
-   * <p>This helper method is for internal and testing usage only.
+   * Add a FQN (fully-qualified name) to Lineage. Segments will be processed via {@link
+   * #getFQNParts}.
    */
-  @Internal
-  public static String getFqName(String system, Iterable<String> segments) {
-    return getFqName(system, null, segments);
+  public void add(
+      String system,
+      @Nullable String subtype,
+      Iterable<String> segments,
+      @Nullable String lastSegmentSep) {
+    add(getFQNParts(system, subtype, segments, lastSegmentSep));
   }
 
   /**
-   * Add a FQN (fully-qualified name) to Lineage. Segments will be processed via {@link #getFqName}.
+   * Add a FQN (fully-qualified name) to Lineage. Segments will be processed via {@link
+   * #getFQNParts}.
    */
-  public void add(String system, @Nullable String subtype, Iterable<String> segments) {
-    add(getFqName(system, subtype, segments));
+  public void add(String system, Iterable<String> segments, @Nullable String lastSegmentSep) {
+    add(system, null, segments, lastSegmentSep);
   }
 
   /**
-   * Add a FQN (fully-qualified name) to Lineage. Segments will be processed via {@link #getFqName}.
+   * Add a FQN (fully-qualified name) to Lineage. Segments will be processed via {@link
+   * #getFQNParts}.
    */
   public void add(String system, Iterable<String> segments) {
-    add(system, null, segments);
+    add(system, segments, null);
   }
 
   /**
-   * Adds the given details as Lineage. For asset level lineage the resource location should be
-   * specified as Dataplex FQN https://cloud.google.com/data-catalog/docs/fully-qualified-names
+   * Adds the given fqn as lineage.
+   *
+   * @param rollupSegments should be an iterable of strings whose concatenation is a valid <a
+   *     href="https://cloud.google.com/data-catalog/docs/fully-qualified-names">Dataplex FQN </a>
+   *     which is already escaped.
+   *     <p>In particular, this means they will often have trailing delimiters.
    */
-  public void add(String details) {
-    metric.add(details);
+  public void add(Iterable<String> rollupSegments) {
+    ImmutableList<String> segments = ImmutableList.copyOf(rollupSegments);
+    this.metric.add(segments);
   }
 
-  /** Query {@link StringSet} metrics from {@link MetricResults}. */
-  public static Set<String> query(MetricResults results, Type type) {
+  /**
+   * Query {@link BoundedTrie} metrics from {@link MetricResults}.
+   *
+   * @param results FQNs from the result.
+   * @param type sources or sinks.
+   * @param truncatedMarker the marker to use to represent truncated FQNs.
+   * @return A flat representation of all FQNs. If the FQN was truncated then it has a trailing
+   *     truncatedMarker.
+   */
+  public static Set<String> query(MetricResults results, Type type, String truncatedMarker) {
     MetricsFilter filter =
         MetricsFilter.builder()
             .addNameFilter(MetricNameFilter.named(LINEAGE_NAMESPACE, type.toString()))
             .build();
     Set<String> result = new HashSet<>();
-    for (MetricResult<StringSetResult> metrics : results.queryMetrics(filter).getStringSets()) {
+    truncatedMarker = truncatedMarker == null ? "*" : truncatedMarker;
+    for (MetricResult<BoundedTrieResult> metrics : results.queryMetrics(filter).getBoundedTries()) {
       try {
-        result.addAll(metrics.getCommitted().getStringSet());
+        for (List<String> fqn : metrics.getCommitted().getResult()) {
+          String end = Boolean.parseBoolean(fqn.get(fqn.size() - 1)) ? truncatedMarker : "";
+          result.add(String.join("", fqn.subList(0, fqn.size() - 1)) + end);
+        }
       } catch (UnsupportedOperationException unused) {
         // MetricsResult.getCommitted throws this exception when runner support missing, just skip.
       }
-      result.addAll(metrics.getAttempted().getStringSet());
+      for (List<String> fqn : metrics.getAttempted().getResult()) {
+        String end = Boolean.parseBoolean(fqn.get(fqn.size() - 1)) ? truncatedMarker : "";
+        result.add(String.join("", fqn.subList(0, fqn.size() - 1)) + end);
+      }
     }
     return result;
   }
 
+  /**
+   * Query {@link BoundedTrie} metrics from {@link MetricResults}.
+   *
+   * @param results FQNs from the result
+   * @param type sources or sinks
+   * @return A flat representation of all FQNs. If the FQN was truncated then it has a trailing '*'.
+   */
+  public static Set<String> query(MetricResults results, Type type) {
+    return query(results, type, "*");
+  }
+
   /** Lineage metrics resource types. */
   public enum Type {
-    SOURCE("sources"),
-    SINK("sinks");
+    SOURCE("sources_v2"),
+    SINK("sinks_v2");
 
     private final String name;
 
@@ -167,4 +197,29 @@ public String toString() {
       return name;
     }
   }
+
+  /**
+   * Wrap segment to valid segment name.
+   *
+   * <p>It escapes reserved characters
+   *
+   * <ul>
+   *   <li>Reserved characters are backtick, colon, whitespace (space, \t, \n) and dot.
+   *   <li>Only segments containing reserved characters must be escaped.
+   *   <li>Segments cannot be escaped partially (i.e. “bigquery:com`.`google.test”).
+   *   <li>Segments must be escaped using backticks (a.k.a. graves).
+   *   <li>Backticks must be escaped using backtick (i.e. bigquery:`test``test`) and the segment
+   *       itself must be escaped as well.
+   * </ul>
+   */
+  @Internal
+  public static String wrapSegment(String value) {
+    value = value.replace("`", "``"); // Escape backticks
+    // the escaped backticks will not throw this off since escaping will
+    // happen if it contains ` in first place.
+    if (RESERVED_CHARS.matcher(value).find()) {
+      return String.format("`%s`", value);
+    }
+    return value;
+  }
 }

sdks/java/core/src/test/java/org/apache/beam/sdk/metrics/LineageTest.java

@@ -20,6 +20,8 @@
 import static org.junit.Assert.assertEquals;
 
 import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.junit.Test;
@@ -30,30 +32,66 @@
 @RunWith(JUnit4.class)
 public class LineageTest {
   @Test
-  public void testGetFqName() {
+  public void testWrapSegment() {
     Map<String, String> testCases =
         ImmutableMap.<String, String>builder()
             .put("apache-beam", "apache-beam")
-            .put("`apache-beam`", "`apache-beam`")
+            .put("`apache-beam`", "```apache-beam```")
             .put("apache.beam", "`apache.beam`")
             .put("apache:beam", "`apache:beam`")
             .put("apache beam", "`apache beam`")
-            .put("`apache beam`", "`apache beam`")
+            .put("apache`beam", "`apache``beam`")
+            .put("apache` beam", "`apache`` beam`")
+            .put("`apache-beam", "```apache-beam`")
+            .put("apache-beam`", "`apache-beam```")
+            .put("`apache beam`", "```apache beam```")
             .put("apache\tbeam", "`apache\tbeam`")
             .put("apache\nbeam", "`apache\nbeam`")
             .build();
     testCases.forEach(
-        (key, value) ->
-            assertEquals("apache:" + value, Lineage.getFqName("apache", ImmutableList.of(key))));
+        (key, value) -> {
+          Iterable<String> fqnPartsIterator =
+              Lineage.getFQNParts("apache", null, ImmutableList.of(key), null);
+          String fqnPartsString = getFqnPartsString(fqnPartsIterator);
+          assertEquals("apache:" + value, fqnPartsString);
+        });
     testCases.forEach(
-        (key, value) ->
-            assertEquals(
-                "apache:beam:" + value,
-                Lineage.getFqName("apache", "beam", ImmutableList.of(key))));
+        (key, value) -> {
+          Iterable<String> fqnPartsIterator =
+              Lineage.getFQNParts("apache", "beam", ImmutableList.of(key), null);
+          String fqnPartsString = getFqnPartsString(fqnPartsIterator);
+          assertEquals("apache:beam:" + value, fqnPartsString);
+        });
     testCases.forEach(
-        (key, value) ->
-            assertEquals(
-                "apache:beam:" + value + "." + value,
-                Lineage.getFqName("apache", "beam", ImmutableList.of(key, key))));
+        (key, value) -> {
+          Iterable<String> fqnPartsIterator =
+              Lineage.getFQNParts("apache", "beam", ImmutableList.of(key, key), null);
+          String fqnPartsString = getFqnPartsString(fqnPartsIterator);
+          assertEquals("apache:beam:" + value + "." + value, fqnPartsString);
+        });
+  }
+
+  @Test
+  public void getFQNParts() {
+    Iterable<String> simpleFQN =
+        Lineage.getFQNParts("system", null, ImmutableList.of("project", "dataset", "table"), null);
+    assertEquals("system:project.dataset.table", getFqnPartsString(simpleFQN));
+
+    Iterable<String> subTypeFQN =
+        Lineage.getFQNParts("system", "topic", ImmutableList.of("project", "topicid"), null);
+    assertEquals("system:topic:project.topicid", getFqnPartsString(subTypeFQN));
+
+    Iterable<String> pathFQN =
+        Lineage.getFQNParts("system", null, ImmutableList.of("bucket", "dir1/dir2/file"), "/");
+    assertEquals("system:bucket.dir1/dir2/file", getFqnPartsString(pathFQN));
+
+    Iterable<String> pathFQNReserved =
+        Lineage.getFQNParts("system", null, ImmutableList.of("bucket", "dir1/dir.2/file"), "/");
+    assertEquals("system:bucket.`dir1/dir.2/file`", getFqnPartsString(pathFQNReserved));
+  }
+
+  private static String getFqnPartsString(Iterable<String> fqnPartsIterable) {
+    return StreamSupport.stream(fqnPartsIterable.spliterator(), false)
+        .collect(Collectors.joining(""));
   }
 }

sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/storage/GcsFileSystem.java

@@ -229,7 +229,7 @@ protected void reportLineage(GcsResourceId resourceId, Lineage lineage, LineageL
       if (level != LineageLevel.TOP_LEVEL && !path.getObject().isEmpty()) {
         segments.add(path.getObject());
       }
-      lineage.add("gcs", segments.build());
+      lineage.add("gcs", segments.build(), "/");
     } else {
       LOG.warn("Report Lineage on relative path {} is unsupported", path.getObject());
     }

sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/storage/GcsFileSystemTest.java

@@ -250,7 +250,7 @@ private void verifyLineage(String uri, List<String> expected) {
     GcsResourceId path = GcsResourceId.fromGcsPath(GcsPath.fromUri(uri));
     Lineage mockLineage = mock(Lineage.class);
     gcsFileSystem.reportLineage(path, mockLineage);
-    verify(mockLineage, times(1)).add("gcs", expected);
+    verify(mockLineage, times(1)).add("gcs", expected, "/");
   }
 
   private StorageObject createStorageObject(String gcsFilename, long fileSize) {

sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/s3/S3FileSystem.java

@@ -668,7 +668,7 @@ protected void reportLineage(S3ResourceId resourceId, Lineage lineage, LineageLe
     if (level != LineageLevel.TOP_LEVEL && !resourceId.getKey().isEmpty()) {
       segments.add(resourceId.getKey());
     }
-    lineage.add("s3", segments.build());
+    lineage.add("s3", segments.build(), "/");
   }
 
   /**