foo

Author: reuvenlax

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java

@@ -222,6 +222,7 @@ private DoFnSignatures() {}
   private static final Collection<Class<? extends Parameter>>
       ALLOWED_ON_WINDOW_EXPIRATION_PARAMETERS =
           ImmutableList.of(
+              Parameter.OnWindowExpirationContextParameter.class,
               Parameter.WindowParameter.class,
               Parameter.PipelineOptionsParameter.class,
               Parameter.OutputReceiverParameter.class,

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java

@@ -194,6 +194,7 @@ public void process(
 
     @OnTimer("retryMismatchedRowsTimer")
     public void onTimer(
+        OnTimerContext context,
         @Key ShardedKey<DestinationT> shardedDestination,
         @Timestamp Instant timestamp,
         @StateId("mismatchedRows") BagState<MismatchedRow> mismatchedRowsBag,
@@ -203,7 +204,7 @@ public void onTimer(
         PipelineOptions pipelineOptions,
         MultiOutputReceiver o)
         throws Exception {
-      // TODO XXX SET DYNAMIC DESTINATIONS
+      dynamicDestinations.setSideInputAccessorFromOnTimerContext(context);
 
       mismatchedRowsBag.readLater();
       currentTimerValue.readLater();

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java

@@ -82,6 +82,33 @@ interface SideInputAccessor {
   private transient @Nullable SideInputAccessor sideInputAccessor;
   private transient @Nullable PipelineOptions options;
 
+  static class SideInputAccessorViaOnTimerContext implements SideInputAccessor {
+    private DoFn<?, ?>.OnTimerContext onTimerContext;
+
+    public SideInputAccessorViaOnTimerContext(DoFn<?, ?>.OnTimerContext onTimerContext) {
+      this.onTimerContext = onTimerContext;
+    }
+
+    @Override
+    public <SideInputT> SideInputT sideInput(PCollectionView<SideInputT> view) {
+      return onTimerContext.sideInput(view);
+    }
+  }
+
+  static class SideInputAccessorViaOnWindowExpirationContext implements SideInputAccessor {
+    private DoFn<?, ?>.OnWindowExpirationContext onWindowExpirationContext;
+
+    public SideInputAccessorViaOnWindowExpirationContext(
+        DoFn<?, ?>.OnWindowExpirationContext onWindowExpirationContext) {
+      this.onWindowExpirationContext = onWindowExpirationContext;
+    }
+
+    @Override
+    public <SideInputT> SideInputT sideInput(PCollectionView<SideInputT> view) {
+      return onWindowExpirationContext.sideInput(view);
+    }
+  }
+
   static class SideInputAccessorViaProcessContext implements SideInputAccessor {
     private DoFn<?, ?>.ProcessContext processContext;
 
@@ -129,6 +156,17 @@ void setSideInputAccessorFromProcessContext(DoFn<?, ?>.ProcessContext context) {
     this.options = context.getPipelineOptions();
   }
 
+  void setSideInputAccessorFromOnTimerContext(DoFn<?, ?>.OnTimerContext context) {
+    this.sideInputAccessor = new SideInputAccessorViaOnTimerContext(context);
+    this.options = context.getPipelineOptions();
+  }
+
+  void setSideInputAccessorFromOnWindowExpirationContext(
+      DoFn<?, ?>.OnWindowExpirationContext context) {
+    this.sideInputAccessor = new SideInputAccessorViaOnWindowExpirationContext(context);
+    this.options = context.getPipelineOptions();
+  }
+
   /**
    * Returns an object that represents at a high level which table is being written to. May not
    * return null.

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java

@@ -215,6 +215,19 @@ void setSideInputAccessorFromProcessContext(DoFn<?, ?>.ProcessContext context) {
       inner.setSideInputAccessorFromProcessContext(context);
     }
 
+    @Override
+    void setSideInputAccessorFromOnTimerContext(DoFn<?, ?>.OnTimerContext context) {
+      super.setSideInputAccessorFromOnTimerContext(context);
+      inner.setSideInputAccessorFromOnTimerContext(context);
+    }
+
+    @Override
+    void setSideInputAccessorFromOnWindowExpirationContext(
+        DoFn<?, ?>.OnWindowExpirationContext context) {
+      super.setSideInputAccessorFromOnWindowExpirationContext(context);
+      inner.setSideInputAccessorFromOnWindowExpirationContext(context);
+    }
+
     @Override
     public String toString() {
       return MoreObjects.toStringHelper(this).add("inner", inner).toString();

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java

@@ -165,6 +165,7 @@ public Duration getAllowedTimestampSkew() {
 
   @OnTimer("pollTimer")
   public void onPollTimer(
+      OnTimerContext context,
       @Key ShardedKey<DestinationT> key,
       PipelineOptions pipelineOptions,
       @StateId("bufferedElements") BagState<TimestampedValue<ElementT>> bag,
@@ -174,6 +175,8 @@ public void onPollTimer(
       BoundedWindow window,
       MultiOutputReceiver o)
       throws Exception {
+    convertMessagesDoFn.getDynamicDestinations().setSideInputAccessorFromOnTimerContext(context);
+
     if (tryFlushBuffer(key.getKey(), pipelineOptions, bag, minBufferedTimestamp, o)) {
       timerTs.clear();
     } else {
@@ -188,12 +191,19 @@ public void onPollTimer(
 
   @OnWindowExpiration
   public void onWindowExpiration(
+      //   OnWindowExpirationContext context,
       @Key ShardedKey<DestinationT> key,
       PipelineOptions pipelineOptions,
       @StateId("bufferedElements") BagState<TimestampedValue<ElementT>> bag,
       @StateId("minBufferedTimestamp") CombiningState<Long, long[], Long> minBufferedTimestamp,
       MultiOutputReceiver o)
       throws Exception {
+    // TODO: Beam doesn't current support OnWindowExpirationContext. Reenable after this is fixed.
+    // https://github.com/apache/beam/issues/38875
+    // convertMessagesDoFn
+    //     .getDynamicDestinations()
+    //     .setSideInputAccessorFromOnWindowExpirationContext(context);
+
     // This can happen on test completion or drain. We can't set any more timers in window
     // expiration, so we just have to loop until the schema is updated.
     BackOff backoff =

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinations.java

@@ -61,4 +61,17 @@ void setSideInputAccessorFromProcessContext(DoFn<?, ?>.ProcessContext context) {
     super.setSideInputAccessorFromProcessContext(context);
     inner.setSideInputAccessorFromProcessContext(context);
   }
+
+  @Override
+  void setSideInputAccessorFromOnTimerContext(DoFn<?, ?>.OnTimerContext context) {
+    super.setSideInputAccessorFromOnTimerContext(context);
+    inner.setSideInputAccessorFromOnTimerContext(context);
+  }
+
+  @Override
+  void setSideInputAccessorFromOnWindowExpirationContext(
+      DoFn<?, ?>.OnWindowExpirationContext context) {
+    super.setSideInputAccessorFromOnWindowExpirationContext(context);
+    inner.setSideInputAccessorFromOnWindowExpirationContext(context);
+  }
 }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java

@@ -1238,6 +1238,7 @@ private void processPayloads(
 
     @OnTimer("retryMismatchedRowsTimer")
     public void onMismatchedRowsTimer(
+        OnTimerContext context,
         PipelineOptions pipelineOptions,
         @Key ShardedKey<DestinationT> shardedDestination,
         @Timestamp org.joda.time.Instant elementTs,
@@ -1251,11 +1252,12 @@ public void onMismatchedRowsTimer(
         @StateId("currentMismatchedRowTimerValue") ValueState<Long> currentTimerValue,
         @StateId("minPendingTimestamp") ValueState<Long> minPendingTimestamp)
         throws Exception {
+      dynamicDestinations.setSideInputAccessorFromOnTimerContext(context);
+
       mismatchedRowsBag.readLater();
       currentTimerValue.readLater();
       minPendingTimestamp.readLater();
 
-      // TODOTDO XXX add context for side inputs.
       TableDestination tableDestination =
           destinations.computeIfAbsent(
               shardedDestination.getKey(),