diff --git a/core/src/main/resources/db/migration/V2.5__Fix_Subscription_MIgration.sql b/core/src/main/resources/db/migration/V2.5__Fix_Subscription_MIgration.sql new file mode 100644 index 00000000000..4ef3013c2c7 --- /dev/null +++ b/core/src/main/resources/db/migration/V2.5__Fix_Subscription_MIgration.sql @@ -0,0 +1,14 @@ +WITH updates as ( + select name, + array_to_string(array_agg(b.array_to_string), ',') as subscriptions + from ( + select name, array_to_string(string_to_array(subscriptions, ':') || '{false}', ':') + from ( + select name, unnest(string_to_array(subscriptions, ',')) as subscriptions from stores) a + WHERE array_length(string_to_array(subscriptions, ':'), 1) = 2) b + group by name +) +UPDATE stores +SET subscriptions = updates.subscriptions +FROM updates +WHERE stores.name = updates.name \ No newline at end of file diff --git a/storage/connectors/bigquery/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadsWithResult.java b/storage/connectors/bigquery/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadsWithResult.java index c0c3f6f84c2..764bfc1f54d 100644 --- a/storage/connectors/bigquery/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadsWithResult.java +++ b/storage/connectors/bigquery/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadsWithResult.java @@ -13,9 +13,7 @@ import org.apache.beam.sdk.coders.*; import org.apache.beam.sdk.options.ValueProvider; -import org.apache.beam.sdk.state.StateSpec; -import org.apache.beam.sdk.state.StateSpecs; -import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.state.*; import org.apache.beam.sdk.transforms.*; import org.apache.beam.sdk.transforms.windowing.*; import org.apache.beam.sdk.values.*; @@ -232,20 +230,28 @@ public void process(ProcessContext c) { "CreateJobId", ParDo.of( new DoFn, String>() { - @StateId("generatedForWindow") - private final StateSpec> generatedForWindow = StateSpecs.value(BooleanCoder.of()); + @StateId("oncePerWindow") + private final StateSpec> oncePerWindow = StateSpecs.set(BooleanCoder.of()); @ProcessElement public void process( ProcessContext c, BoundedWindow w, - @StateId("generatedForWindow") ValueState generatedForWindow) { + @StateId("oncePerWindow") SetState oncePerWindow) { - if (generatedForWindow.read() != null) { + // if set already contains something + // it means we already generated Id for this window + Boolean empty = oncePerWindow.isEmpty().read(); + if (empty != null && !empty) { return; } - generatedForWindow.write(true); + // trying to add to Set and check if it was added + // if true - we won and Id will be generated in current Process + Boolean insertResult = oncePerWindow.addIfAbsent(true).read(); + if (insertResult != null && !insertResult) { + return; + } c.output( String.format( diff --git a/tests/e2e/redis/basic-ingest-redis-serving.py b/tests/e2e/redis/basic-ingest-redis-serving.py index 16b4b1f41d1..5bb79bc8e8f 100644 --- a/tests/e2e/redis/basic-ingest-redis-serving.py +++ b/tests/e2e/redis/basic-ingest-redis-serving.py @@ -558,32 +558,6 @@ def try_get_features2(): ) -@pytest.mark.timeout(300) -@pytest.mark.run(order=19) -def test_basic_ingest_jobs(client): - # list ingestion jobs given featureset - cust_trans_fs = client.get_feature_set(name="customer_transactions") - ingest_jobs = client.list_ingest_jobs( - feature_set_ref=FeatureSetRef.from_feature_set(cust_trans_fs) - ) - # filter ingestion jobs to only those that are running - ingest_jobs = [ - job for job in ingest_jobs if job.status == IngestionJobStatus.RUNNING - ] - assert len(ingest_jobs) >= 1 - - for ingest_job in ingest_jobs: - # restart ingestion ingest_job - client.restart_ingest_job(ingest_job) - ingest_job.wait(IngestionJobStatus.RUNNING) - assert ingest_job.status == IngestionJobStatus.RUNNING - - # stop ingestion ingest_job - client.stop_ingest_job(ingest_job) - ingest_job.wait(IngestionJobStatus.ABORTED) - assert ingest_job.status == IngestionJobStatus.ABORTED - - @pytest.fixture(scope="module") def all_types_dataframe(): return pd.DataFrame( @@ -762,16 +736,33 @@ def test_all_types_ingest_jobs(client, all_types_dataframe): ] assert len(ingest_jobs) >= 1 - for ingest_job in ingest_jobs: - # restart ingestion ingest_job - client.restart_ingest_job(ingest_job) - ingest_job.wait(IngestionJobStatus.RUNNING) - assert ingest_job.status == IngestionJobStatus.RUNNING + ingest_job = ingest_jobs[0] + # restart ingestion ingest_job + # restart means stop current job + # (replacement will be automatically spawned) + client.restart_ingest_job(ingest_job) + # wait for replacement to be created + time.sleep(15) # should be more than polling_interval + + # id without timestamp part + # that remains the same between jobs + shared_id = "-".join(ingest_job.id.split("-")[:-1]) + replacement_jobs = [ + job + for job in ingest_jobs + if job.status == IngestionJobStatus.RUNNING and job.id.startswith(shared_id) + ] - # stop ingestion ingest_job - client.stop_ingest_job(ingest_job) - ingest_job.wait(IngestionJobStatus.ABORTED) - assert ingest_job.status == IngestionJobStatus.ABORTED + assert len(replacement_jobs) >= 1 + replacement_job = replacement_jobs[0] + + replacement_job.wait(IngestionJobStatus.RUNNING) + assert replacement_job.status == IngestionJobStatus.RUNNING + + # stop ingestion ingest_job + client.stop_ingest_job(replacement_job) + replacement_job.wait(IngestionJobStatus.ABORTED) + assert replacement_job.status == IngestionJobStatus.ABORTED @pytest.fixture(scope="module")