diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json index 2e255c8f3cf6..2bb052d5f715 100644 --- a/.github/trigger_files/beam_PostCommit_Python.json +++ b/.github/trigger_files/beam_PostCommit_Python.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", "pr": "37345", - "modification": 52 + "modification": 53 } diff --git a/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json b/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json index b5704c67ef1c..86bf1193abd9 100644 --- a/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json +++ b/.github/trigger_files/beam_PostCommit_Yaml_Xlang_Direct.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "revision": 6 + "revision": 7 } diff --git a/.github/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml b/.github/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml index 6d34bd1d156f..7efb81229364 100644 --- a/.github/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml +++ b/.github/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml @@ -41,7 +41,7 @@ jobs: runs-on: [self-hosted, ubuntu-24.04, main] timeout-minutes: 30 steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Setup gcloud uses: google-github-actions/setup-gcloud@aa5489c8933f4cc7a4f7d45035b3b1440c9c10db diff --git a/.github/workflows/beam_PreCommit_GHA.yml b/.github/workflows/beam_PreCommit_GHA.yml index af5513a2835e..ae562e01934d 100644 --- a/.github/workflows/beam_PreCommit_GHA.yml +++ b/.github/workflows/beam_PreCommit_GHA.yml @@ -42,7 +42,7 @@ permissions: packages: read pages: read repository-projects: read - security-events: read + security-events: write statuses: read # This allows a subsequently queued workflow run to interrupt previous runs @@ -78,6 +78,10 @@ jobs: comment_phrase: ${{ matrix.job_phrase }} github_token: ${{ secrets.GITHUB_TOKEN }} github_job: ${{ matrix.job_name }} (${{ matrix.job_phrase }}) + - name: Run zizmor + uses: zizmorcore/zizmor-action@192e21d79ab29983730a13d1382995c2307fbcaa # v0.5.7 + with: + advanced-security: true - name: Setup environment uses: ./.github/actions/setup-environment-action with: diff --git a/sdks/go.mod b/sdks/go.mod index 272f91451db8..7924367a3ef2 100644 --- a/sdks/go.mod +++ b/sdks/go.mod @@ -31,13 +31,13 @@ require ( cloud.google.com/go/profiler v0.6.0 cloud.google.com/go/pubsub v1.50.2 cloud.google.com/go/spanner v1.92.0 - cloud.google.com/go/storage v1.62.3 + cloud.google.com/go/storage v1.63.0 github.com/aws/aws-sdk-go-v2 v1.42.0 github.com/aws/aws-sdk-go-v2/config v1.32.25 github.com/aws/aws-sdk-go-v2/credentials v1.19.24 github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.28 github.com/aws/aws-sdk-go-v2/service/s3 v1.104.0 - github.com/aws/smithy-go v1.27.2 + github.com/aws/smithy-go v1.27.3 github.com/docker/go-connections v0.7.0 // indirect github.com/dustin/go-humanize v1.0.1 github.com/go-sql-driver/mysql v1.10.0 @@ -203,6 +203,6 @@ require ( golang.org/x/mod v0.36.0 // indirect golang.org/x/tools v0.45.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect - google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20260610212136-7ab31c22f7ad // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20260615183401-62b3387ff324 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260615183401-62b3387ff324 // indirect ) diff --git a/sdks/go.sum b/sdks/go.sum index a8692d22a9e9..75ca43710f28 100644 --- a/sdks/go.sum +++ b/sdks/go.sum @@ -99,8 +99,8 @@ cloud.google.com/go/storage v1.8.0/go.mod h1:Wv1Oy7z6Yz3DshWRJFhqM/UCfaWIRTdp0RX cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9ullr3+Kg0= cloud.google.com/go/storage v1.12.0/go.mod h1:fFLk2dp2oAhDz8QFKwqrjdJvxSp/W2g7nillojlL5Ho= cloud.google.com/go/storage v1.21.0/go.mod h1:XmRlxkgPjlBONznT2dDUU/5XlpU2OjMnKuqnZI01LAA= -cloud.google.com/go/storage v1.62.3 h1:SZq1t23NCI+e96dH77Dg3PEfsNNEjqO8zE5AnD8gVD0= -cloud.google.com/go/storage v1.62.3/go.mod h1:cpYz/kRVZ+UQAF1uHeea10/9ewcRbxGoGNKsS9daSXA= +cloud.google.com/go/storage v1.63.0 h1:hvXF2xfg9I32bjujggxgkEZn/Ej6sJ9pieFgeueBLrQ= +cloud.google.com/go/storage v1.63.0/go.mod h1:tirWVptrFNo5GEX2DQ47JooF7yaweJdAJ1hYAVMvKzE= cloud.google.com/go/trace v1.0.0/go.mod h1:4iErSByzxkyHWzzlAj63/Gmjz0NH1ASqhJguHpGcr6A= cloud.google.com/go/trace v1.2.0/go.mod h1:Wc8y/uYyOhPy12KEnXG9XGrvfMz5F5SrYecQlbW1rwM= cloud.google.com/go/trace v1.16.0 h1:GmQovzFc5F0CNfl0VLgL64aoTtu7xsM0YajW2GlG9+E= @@ -274,8 +274,8 @@ github.com/aws/aws-sdk-go-v2/service/sts v1.43.3 h1:VrIhKRCSK1umelSgB9RghvA9RTUY github.com/aws/aws-sdk-go-v2/service/sts v1.43.3/go.mod h1:r8wkDOuLaaMFqFiYAb8dGY2A3gJCOujMc6CFOVC4Zhc= github.com/aws/smithy-go v1.11.2/go.mod h1:3xHYmszWVx2c0kIwQeEVf9uSm4fYZt67FBJnwub1bgM= github.com/aws/smithy-go v1.17.0/go.mod h1:NukqUGpCZIILqqiV0NIjeFh24kd/FAa4beRb6nbIUPE= -github.com/aws/smithy-go v1.27.2 h1:y9NPmSE6am6LjEFPfqHqG/jJk7AauQvhCJONKh7kpzk= -github.com/aws/smithy-go v1.27.2/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= +github.com/aws/smithy-go v1.27.3 h1:F3Zb497UhhskkfpJmfkXswyo+t0sh9OTBnIHjogWbVY= +github.com/aws/smithy-go v1.27.3/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= github.com/bobg/gcsobj v0.1.2/go.mod h1:vS49EQ1A1Ib8FgrL58C8xXYZyOCR2TgzAdopy6/ipa8= github.com/boombuler/barcode v1.0.0/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8= @@ -869,8 +869,8 @@ go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.68.0 h1:CqXxU8V go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.68.0/go.mod h1:BuhAPThV8PBHBvg8ZzZ/Ok3idOdhWIodywz2xEcRbJo= go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= -go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.43.0 h1:TC+BewnDpeiAmcscXbGMfxkO+mwYUwE/VySwvw88PfA= -go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.43.0/go.mod h1:J/ZyF4vfPwsSr9xJSPyQ4LqtcTPULFR64KwTikGLe+A= +go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.44.0 h1:hqxVTu/GtBF+vJ8d1fzW7fRxZFvgoDjWcxwwCaFDYpU= +go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.44.0/go.mod h1:z5fVEF4X5v0ESvlJqBrrFlBVoj5EQuefZpzsu7R+x5Q= go.opentelemetry.io/otel/metric v1.44.0 h1:1w0gILTcHdr3YI+ixLyjemwrVnsMURbTZFrSYCdDdmc= go.opentelemetry.io/otel/metric v1.44.0/go.mod h1:8O7hanEPBNgEMmybD3s2VBKcgWOCsA6tzHBPODAiquo= go.opentelemetry.io/otel/metric/x v0.66.0 h1:YkCrx1zLOChi9ZcZ6euupOcsgzbVlec7D/xoEU1+cTA= @@ -1423,10 +1423,10 @@ google.golang.org/genproto v0.0.0-20220324131243-acbaeb5b85eb/go.mod h1:hAL49I2I google.golang.org/genproto v0.0.0-20220401170504-314d38edb7de/go.mod h1:8w6bsBMX6yCPbAVTeqQHvzxW0EIFigd5lZyahWgyfDo= google.golang.org/genproto v0.0.0-20260523011958-0a33c5d7ca68 h1:cTHF8xtqtBN5sQ4dcoNwOS6FFejvFTkWQbZXsTU3trM= google.golang.org/genproto v0.0.0-20260523011958-0a33c5d7ca68/go.mod h1:RRHjglSYABVCWpQ7USCpdfhcd9t4PkajvVwyynZizTc= -google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa h1:Kjn0N0tCrDgiAFW+lGO4JZ3ck44CehvJQMAwj9QF0G8= -google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:q4lMZS6kskjT5HvCPrnnypcDPVJqT/f4nfxmkE7gryY= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260610212136-7ab31c22f7ad h1:45WmJvIV6C2+O/jjLkPUH+F3aOj/1miDoU2DD0+NWbg= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260610212136-7ab31c22f7ad/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/genproto/googleapis/api v0.0.0-20260615183401-62b3387ff324 h1:g0RAkxK/smSu/iRwC/KIX1mwUoVJtk2OjbgaeS4DmUM= +google.golang.org/genproto/googleapis/api v0.0.0-20260615183401-62b3387ff324/go.mod h1:Z4WJ5pJOYWFWcHEQUelD5QaZDknIQkpIL/+fyJOT9+A= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260615183401-62b3387ff324 h1:9HZDLIdYBJXAnaFOr9WHrKVycfpY+75s9HGadC0305A= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260615183401-62b3387ff324/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= google.golang.org/grpc v1.20.1/go.mod h1:10oTOabMzJvdu6/UiuZezV6QK5dSlG84ov/aaiqXj38= google.golang.org/grpc v1.21.1/go.mod h1:oYelfM1adQP15Ek0mdvEgi9Df8B9CZIaU1084ijfRaM= diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java index f6272ba4fe67..23372be2a2dc 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AsyncWrapper.java @@ -644,15 +644,17 @@ void commitFinishedItems( if (activeElements.containsKey(elementId)) { InFlightElement inFlight = activeElements.get(elementId); + // Future is either completed, cancelled, or throws an exception if (inFlight.future.isDone()) { + // Remove from local active map before checking the result + activeElements.remove(elementId); try { if (!inFlight.future.isCancelled()) { toReturn.add(inFlight.future.get()); + // Only mark as finished if future was not cancelled + finishedElementIds.add(elementId); + itemsFinished++; } - - finishedElementIds.add(elementId); - activeElements.remove(elementId); - itemsFinished++; } catch (Exception e) { LOG.error("Error executing async task for element {}", element, e); throw new RuntimeException("Error executing async task for element " + element, e); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java index 183b1851459c..d098104a8cf6 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/AsyncWrapperTest.java @@ -19,6 +19,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertThrows; import java.io.Serializable; import java.util.ArrayList; @@ -179,7 +180,7 @@ public BagState readLater() { // 4. Used for testing Timer mock implementations. private static class FakeTimer implements Timer { - private Instant time = Instant.EPOCH; + private volatile Instant time = Instant.EPOCH; @Override public void set(Instant absoluteTime) { @@ -874,4 +875,67 @@ public void testResetStateConcurrentTeardown() { // Verify calling resetState() while background tasks are running finishes cleanly AsyncWrapper.resetState(); } + + // Test 15: testTransientRpcFailureRetry + // Verify DoFn exceptions wipe local active state so bundle retries reschedule work. + @Test + public void testTransientRpcFailureRetry() { + FlakyDoFn dofn = new FlakyDoFn(); + AsyncWrapper asyncWrapper = + new AsyncWrapper<>( + dofn, 1, Duration.standardSeconds(5), null, null, null, null, useThreadPool); + asyncWrapper.setup(null); + + FakeBagState> fakeBagState = new FakeBagState<>(); + FakeTimer fakeTimer = new FakeTimer(); + KV msg = KV.of("key1", "1"); + + asyncWrapper.processDirect(msg, GlobalWindow.INSTANCE, Instant.now(), fakeBagState, fakeTimer); + waitForEmpty(asyncWrapper); + + // Attempt 1 should throw RuntimeException wrapping ExecutionException + assertThrows( + RuntimeException.class, + () -> + asyncWrapper.commitFinishedItemsDirect( + fakeTimer.getCurrentRelativeTime(), fakeBagState, fakeTimer)); + + // Verify the failed Future was removed from local active elements + assertEquals(0, asyncWrapper.getItemsInBufferCount()); + + // Simulate runner bundle retry: commitFinishedItemsDirect runs again with msg still in state. + // Because the dead future was removed, it will reschedule msg and succeed on Attempt 2. + waitForEmpty(asyncWrapper); + List result = + asyncWrapper.commitFinishedItemsDirect( + fakeTimer.getCurrentRelativeTime(), fakeBagState, fakeTimer); + if (result.isEmpty()) { + waitForEmpty(asyncWrapper); + result = + asyncWrapper.commitFinishedItemsDirect( + fakeTimer.getCurrentRelativeTime(), fakeBagState, fakeTimer); + } + + checkOutput(result, Collections.singletonList("1")); + assertEquals(0, fakeBagState.items.size()); + } + + private static class FlakyDoFn extends DoFn { + private int attempts = 0; + private final ReentrantLock lock = new ReentrantLock(); + + @ProcessElement + public void processElement(@Element String element, OutputReceiver receiver) { + lock.lock(); + try { + attempts++; + if (attempts == 1) { + throw new RuntimeException("Transient RPC Error"); + } + } finally { + lock.unlock(); + } + receiver.output(element); + } + } } diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py b/sdks/python/apache_beam/io/gcp/gcsio.py index 7cc68fd1e4a2..179525de2cb0 100644 --- a/sdks/python/apache_beam/io/gcp/gcsio.py +++ b/sdks/python/apache_beam/io/gcp/gcsio.py @@ -35,6 +35,7 @@ from typing import Optional from typing import Union +from google.api_core.exceptions import Conflict from google.api_core.exceptions import RetryError from google.cloud import storage from google.cloud.exceptions import NotFound @@ -143,7 +144,17 @@ def get_or_create_default_gcs_bucket(options): 'Creating default GCS bucket for project %s: gs://%s', project, bucket_name) - return gcs.create_bucket(bucket_name, project, location=region) + try: + return gcs.create_bucket(bucket_name, project, location=region) + except Conflict: + bucket = gcs.get_bucket(bucket_name) + if bucket: + _validate_bucket_project( + bucket, + project, + credentials=getattr(gcs.client, '_credentials', None)) + return bucket + raise def create_storage_client(pipeline_options, use_credentials=True): diff --git a/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py b/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py index f5da9b60dbd6..a68b1179faf0 100644 --- a/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py +++ b/sdks/python/apache_beam/io/gcp/gcsio_integration_test.py @@ -33,6 +33,7 @@ import unittest import uuid import zlib +from hashlib import blake2b import mock import pytest @@ -207,19 +208,17 @@ def test_create_default_bucket(self, mock_default_gcs_bucket_name): # requires this option unset. google_cloud_options.dataflow_kms_key = None - import random - from hashlib import blake2b - - # Add a random number to avoid collision if multiple test instances - # are run at the same time. To avoid too many dangling buckets if bucket - # removal fails, we limit the max number of possible bucket names in this - # test to 1000. - overridden_bucket_name = 'gcsio-it-%d-%s-%s-%d' % ( - random.randint(0, 999), + # Add a unique uuid and the parameterized test options to the bucket name + # to avoid collisions when multiple parameterized instances run in parallel + # or concurrent CI jobs run at the same time. + overridden_bucket_name = 'gcsio-it-%s-%s-%s-%d-%s-%s' % ( + uuid.uuid4().hex[:6], google_cloud_options.region, blake2b(google_cloud_options.project.encode('utf8'), - digest_size=4).hexdigest(), - int(time.time())) + digest_size=2).hexdigest(), + int(time.time()), + '1' if self.no_gcsio_throttling_counter else '0', + '1' if self.enable_gcsio_blob_generation else '0') mock_default_gcs_bucket_name.return_value = overridden_bucket_name diff --git a/sdks/python/apache_beam/io/gcp/gcsio_test.py b/sdks/python/apache_beam/io/gcp/gcsio_test.py index 9c4414175e48..ddb94f9c17fa 100644 --- a/sdks/python/apache_beam/io/gcp/gcsio_test.py +++ b/sdks/python/apache_beam/io/gcp/gcsio_test.py @@ -487,6 +487,46 @@ def test_get_or_create_default_gcs_bucket_ownership_mock_project_number( self.assertEqual(bucket, mock_bucket) mock_crm_class.assert_not_called() + @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient') + @mock.patch('apache_beam.io.gcp.gcsio.GcsIO') + def test_get_or_create_default_gcs_bucket_conflict( + self, mock_gcsio_class, mock_crm_class): + mock_gcsio = mock_gcsio_class.return_value + mock_bucket = mock.Mock() + mock_bucket.project_number = 123456789 + mock_gcsio.get_bucket.side_effect = [None, mock_bucket] + + from google.api_core.exceptions import Conflict + mock_gcsio.create_bucket.side_effect = Conflict("Already owned by you") + + mock_crm_client = mock_crm_class.return_value + mock_project_info = mock.Mock() + mock_project_info.name = 'projects/123456789' + mock_crm_client.get_project.return_value = mock_project_info + + options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1') + bucket = gcsio.get_or_create_default_gcs_bucket(options) + + self.assertEqual(bucket, mock_bucket) + self.assertEqual(mock_gcsio.get_bucket.call_count, 2) + mock_gcsio.create_bucket.assert_called_once() + + @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient') + @mock.patch('apache_beam.io.gcp.gcsio.GcsIO') + def test_get_or_create_default_gcs_bucket_conflict_reraise( + self, mock_gcsio_class, mock_crm_class): + mock_gcsio = mock_gcsio_class.return_value + mock_gcsio.get_bucket.side_effect = [None, None] + + from google.api_core.exceptions import Conflict + mock_gcsio.create_bucket.side_effect = Conflict("Bucket name unavailable") + + options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1') + with self.assertRaises(Conflict): + gcsio.get_or_create_default_gcs_bucket(options) + + self.assertEqual(mock_gcsio.get_bucket.call_count, 2) + def test_exists(self): file_name = 'gs://gcsio-test/dummy_file' file_size = 1234 diff --git a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py index e67c5f2a3708..12b200764578 100644 --- a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py +++ b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py @@ -362,6 +362,8 @@ def test_batch_write_with_ordering_key(self): # Retry pulling to handle PubSub delivery delays received_messages = [] + received_message_ids = set() + ack_ids = [] deadline = time.time() + 60 # wait up to 60 seconds while time.time() < deadline: response = self.sub_client.pull( @@ -369,7 +371,13 @@ def test_batch_write_with_ordering_key(self): 'subscription': ordering_sub.name, 'max_messages': 10, }) - received_messages.extend(response.received_messages) + for msg in response.received_messages: + ack_ids.append(msg.ack_id) + # Pub/Sub guarantees at-least-once delivery, so we must deduplicate + # messages by message_id to handle potential duplicate deliveries. + if msg.message.message_id not in received_message_ids: + received_message_ids.add(msg.message.message_id) + received_messages.append(msg) if len(received_messages) >= len(test_messages): break time.sleep(5) @@ -384,12 +392,12 @@ def test_batch_write_with_ordering_key(self): self.assertEqual(received_map[b'order_data002'].ordering_key, 'key1') self.assertEqual(received_map[b'order_data003'].ordering_key, 'key2') - ack_ids = [msg.ack_id for msg in received_messages] - self.sub_client.acknowledge( - request={ - 'subscription': ordering_sub.name, - 'ack_ids': ack_ids, - }) + if ack_ids: + self.sub_client.acknowledge( + request={ + 'subscription': ordering_sub.name, + 'ack_ids': ack_ids, + }) finally: self.sub_client.delete_subscription( request={'subscription': ordering_sub.name}) diff --git a/sdks/python/apache_beam/runners/worker/statesampler_fast.pyx b/sdks/python/apache_beam/runners/worker/statesampler_fast.pyx index 7075ef47017d..5a67d6764ea0 100644 --- a/sdks/python/apache_beam/runners/worker/statesampler_fast.pyx +++ b/sdks/python/apache_beam/runners/worker/statesampler_fast.pyx @@ -153,9 +153,14 @@ cdef class StateSampler(object): self.sampling_thread.join() def reset(self): - for state in self.scoped_states_by_index: - (state)._nsecs = 0 - self.started = self.finished = False + with nogil: + pythread.PyThread_acquire_lock(self.lock, pythread.WAIT_LOCK) + try: + for state in self.scoped_states_by_index: + (state)._nsecs = 0 + self.started = self.finished = False + finally: + pythread.PyThread_release_lock(self.lock) cpdef ScopedState current_state(self): return self.current_state_c() diff --git a/sdks/python/apache_beam/transforms/async_dofn.py b/sdks/python/apache_beam/transforms/async_dofn.py index ad3d5bc66469..f104770d8c0c 100644 --- a/sdks/python/apache_beam/transforms/async_dofn.py +++ b/sdks/python/apache_beam/transforms/async_dofn.py @@ -489,9 +489,11 @@ def commit_finished_items( if x_id in processing_elements: _, future = processing_elements[x_id] if future.done(): + # Pop from local active map before checking the result + processing_elements.pop(x_id) to_return.append(future.result()) + # Only mark as finished if result() succeeded without exception finished_items.append(x) - processing_elements.pop(x_id) items_finished += 1 else: items_not_yet_finished += 1 diff --git a/sdks/python/apache_beam/transforms/async_dofn_test.py b/sdks/python/apache_beam/transforms/async_dofn_test.py index 39901d791fb9..f9e07d57be9b 100644 --- a/sdks/python/apache_beam/transforms/async_dofn_test.py +++ b/sdks/python/apache_beam/transforms/async_dofn_test.py @@ -522,6 +522,53 @@ def test_reset_state_concurrent_teardown(self): else: self.assertEqual(p.exitcode, 0) + def test_transient_rpc_failure_retry(self): + # Verify DoFn exceptions wipe local active state so retries reschedule work. + class FlakyDoFn(beam.DoFn): + def __init__(self): + self.attempts = 0 + self.lock = Lock() + + def process(self, element): + with self.lock: + self.attempts += 1 + current_attempt = self.attempts + if current_attempt == 1: + raise RuntimeError("Transient RPC Error") + yield element + + dofn = FlakyDoFn() + async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio) + async_dofn.setup() + fake_bag_state = FakeBagState([]) + fake_timer = FakeTimer(0) + msg = ('key1', 1) + + async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer) + self.wait_for_empty(async_dofn) + + # Attempt 1 should raise the RuntimeError stored in the future + with self.assertRaises(RuntimeError): + async_dofn.commit_finished_items(fake_bag_state, fake_timer) + + # Verify the failed future was popped from local processing_elements + with async_lib.AsyncWrapper._lock: + self.assertNotIn( + async_dofn._id_fn(msg[1]), + async_lib.AsyncWrapper._processing_elements[async_dofn._uuid], + ) + + # Simulate runner bundle retry: commit_finished_items runs again with msg still in state. + # Because the dead future was popped, it will reschedule msg and succeed on Attempt 2. + self.wait_for_empty(async_dofn) + result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) + if not result: + self.wait_for_empty(async_dofn) + result = async_dofn.commit_finished_items(fake_bag_state, fake_timer) + + self.check_output(result, [msg]) + self.assertEqual(fake_bag_state.items, []) + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/yaml/extended_tests/databases/iceberg.yaml b/sdks/python/apache_beam/yaml/extended_tests/databases/iceberg.yaml index d7449233aab5..98462f2d0271 100644 --- a/sdks/python/apache_beam/yaml/extended_tests/databases/iceberg.yaml +++ b/sdks/python/apache_beam/yaml/extended_tests/databases/iceberg.yaml @@ -16,8 +16,18 @@ # fixtures: - - name: TEMP_DIR - type: "tempfile.TemporaryDirectory" + - name: ICEBERG_FIXTURE + type: "apache_beam.yaml.integration_tests.temp_iceberg_table_with_pk" + config: + table_data: + name: "labels" + schema: + type: "struct" + schema-id: 0 + fields: + - { id: 1, name: "label", required: true, type: "string" } + - { id: 2, name: "rank", required: true, type: "long" } + identifier-field-ids: [1] pipelines: - name: write @@ -32,14 +42,14 @@ pipelines: - {label: "389a", rank: 2} - type: WriteToIceberg config: - table: db.labels - catalog_name: hadoop_catalog + table: "{ICEBERG_FIXTURE[table]}" + catalog_name: rest_catalog catalog_properties: - type: hadoop - warehouse: "{TEMP_DIR}" + type: rest + uri: "{ICEBERG_FIXTURE[api_url]}" options: project: "apache-beam-testing" - temp_location: "{TEMP_DIR}" + temp_location: "{ICEBERG_FIXTURE[temp_dir]}" - name: read pipeline: @@ -47,11 +57,11 @@ pipelines: transforms: - type: ReadFromIceberg config: - table: db.labels - catalog_name: hadoop_catalog + table: "{ICEBERG_FIXTURE[table]}" + catalog_name: rest_catalog catalog_properties: - type: hadoop - warehouse: "{TEMP_DIR}" + type: rest + uri: "{ICEBERG_FIXTURE[api_url]}" - type: AssertEqual config: elements: @@ -60,7 +70,7 @@ pipelines: - {label: "389a", rank: 2} options: project: "apache-beam-testing" - temp_location: "{TEMP_DIR}" + temp_location: "{ICEBERG_FIXTURE[temp_dir]}" - name: read_cdc_batch pipeline: @@ -68,11 +78,11 @@ pipelines: transforms: - type: ReadFromIcebergCDC config: - table: db.labels - catalog_name: hadoop_catalog + table: "{ICEBERG_FIXTURE[table]}" + catalog_name: rest_catalog catalog_properties: - type: hadoop - warehouse: "{TEMP_DIR}" + type: rest + uri: "{ICEBERG_FIXTURE[api_url]}" from_timestamp: 1762819200000 to_timestamp: 2078352000000 filter: '"label" = ''11a'' or "rank" = 1' @@ -86,7 +96,7 @@ pipelines: - {label: "37a", rank: 1} options: project: "apache-beam-testing" - temp_location: "{TEMP_DIR}" + temp_location: "{ICEBERG_FIXTURE[temp_dir]}" - name: read_cdc_streaming pipeline: @@ -94,11 +104,11 @@ pipelines: transforms: - type: ReadFromIcebergCDC config: - table: db.labels - catalog_name: hadoop_catalog + table: "{ICEBERG_FIXTURE[table]}" + catalog_name: rest_catalog catalog_properties: - type: hadoop - warehouse: "{TEMP_DIR}" + type: rest + uri: "{ICEBERG_FIXTURE[api_url]}" streaming: True to_timestamp: 2078352000000 filter: '"label" = ''11a'' or "rank" = 1' @@ -112,4 +122,4 @@ pipelines: - {label: "37a", rank: 1} options: project: "apache-beam-testing" - temp_location: "{TEMP_DIR}" + temp_location: "{ICEBERG_FIXTURE[temp_dir]}" diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index e319a3d3a9bd..a7abd223b295 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -26,15 +26,20 @@ import os import random import secrets +import shutil import sqlite3 import string import struct +import tempfile +import time import unittest import uuid from datetime import datetime from datetime import timezone import mock +import requests +from testcontainers.core.container import DockerContainer from apache_beam.coders import Coder from apache_beam.coders.coder_impl import CoderImpl @@ -556,6 +561,75 @@ def temp_oracle_database(): yield f"jdbc:oracle:thin:system/oracle@localhost:{port}/XEPDB1" +@contextlib.contextmanager +def temp_iceberg_table_with_pk(table_data): + + # Create a temp dir that will be shared between host and container. + # We use the exact same path on both to avoid path mapping issues. + # We create it in the current working directory (workspace) because + # Docker in GitHub Actions often cannot mount directories from /tmp. + temp_dir = tempfile.mkdtemp(dir=os.getcwd()) + os.chmod(temp_dir, 0o777) + + # Start the Iceberg REST catalog container + container = DockerContainer("tabulario/iceberg-rest:0.6.0") + container.with_exposed_ports(8181) + container.with_volume_mapping(temp_dir, temp_dir, mode='rw') + container.with_env("HADOOP_USER_NAME", "iceberg") + container.with_env("CATALOG_WAREHOUSE", temp_dir) + container.with_env( + "CATALOG_IO__IMPL", "org.apache.iceberg.hadoop.HadoopFileIO") + + try: + container.start() + + ip = container.get_container_host_ip() + port = container.get_exposed_port(8181) + api_url = f"http://{ip}:{port}" + + # Poll the REST API until it is ready + for _ in range(30): + try: + response = requests.get(f"{api_url}/v1/config", timeout=5) + if response.status_code == 200: + break + except (requests.exceptions.ConnectionError, requests.exceptions.Timeout): + pass + time.sleep(1) + else: + raise RuntimeError("Iceberg REST catalog failed to start in time.") + + # Create namespace 'db' + requests.post( + f"{api_url}/v1/namespaces", + json={"namespace": ["db"]}, + headers={"Content-Type": "application/json"}, + timeout=10) + + # Create table with primary key + response = requests.post( + f"{api_url}/v1/namespaces/db/tables", + json=table_data, + headers={"Content-Type": "application/json"}, + timeout=10) + if response.status_code != 200: + raise RuntimeError(f"Failed to create Iceberg table: {response.text}") + + # Change permissions of the created directories inside the container + # so the host user can write to them. + container.get_wrapped_container().exec_run(f"chmod -R 777 {temp_dir}") + + yield { + "api_url": api_url, + "temp_dir": temp_dir, + "table": f"db.{table_data['name']}" + } + + finally: + container.stop() + shutil.rmtree(temp_dir, ignore_errors=True) + + @contextlib.contextmanager def temp_kafka_server(): """Context manager to provide a temporary Kafka server for testing.