Skip to content

Commit 4c1b152

Browse files
author
Oleksii Moskalenko
authored
Sync v0.6 branch with master (#872)
* correct subscription migration * fix ingestion job restart e2e test * delete duplicated test * bq: fix job id generator race condition
1 parent ee94df0 commit 4c1b152

3 files changed

Lines changed: 54 additions & 43 deletions

File tree

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
WITH updates as (
2+
select name,
3+
array_to_string(array_agg(b.array_to_string), ',') as subscriptions
4+
from (
5+
select name, array_to_string(string_to_array(subscriptions, ':') || '{false}', ':')
6+
from (
7+
select name, unnest(string_to_array(subscriptions, ',')) as subscriptions from stores) a
8+
WHERE array_length(string_to_array(subscriptions, ':'), 1) = 2) b
9+
group by name
10+
)
11+
UPDATE stores
12+
SET subscriptions = updates.subscriptions
13+
FROM updates
14+
WHERE stores.name = updates.name

storage/connectors/bigquery/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadsWithResult.java

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,7 @@
1313

1414
import org.apache.beam.sdk.coders.*;
1515
import org.apache.beam.sdk.options.ValueProvider;
16-
import org.apache.beam.sdk.state.StateSpec;
17-
import org.apache.beam.sdk.state.StateSpecs;
18-
import org.apache.beam.sdk.state.ValueState;
16+
import org.apache.beam.sdk.state.*;
1917
import org.apache.beam.sdk.transforms.*;
2018
import org.apache.beam.sdk.transforms.windowing.*;
2119
import org.apache.beam.sdk.values.*;
@@ -232,20 +230,28 @@ public void process(ProcessContext c) {
232230
"CreateJobId",
233231
ParDo.of(
234232
new DoFn<KV<Void, TableRow>, String>() {
235-
@StateId("generatedForWindow")
236-
private final StateSpec<ValueState<Boolean>> generatedForWindow = StateSpecs.value(BooleanCoder.of());
233+
@StateId("oncePerWindow")
234+
private final StateSpec<SetState<Boolean>> oncePerWindow = StateSpecs.set(BooleanCoder.of());
237235

238236
@ProcessElement
239237
public void process(
240238
ProcessContext c,
241239
BoundedWindow w,
242-
@StateId("generatedForWindow") ValueState<Boolean> generatedForWindow) {
240+
@StateId("oncePerWindow") SetState<Boolean> oncePerWindow) {
243241

244-
if (generatedForWindow.read() != null) {
242+
// if set already contains something
243+
// it means we already generated Id for this window
244+
Boolean empty = oncePerWindow.isEmpty().read();
245+
if (empty != null && !empty) {
245246
return;
246247
}
247248

248-
generatedForWindow.write(true);
249+
// trying to add to Set and check if it was added
250+
// if true - we won and Id will be generated in current Process
251+
Boolean insertResult = oncePerWindow.addIfAbsent(true).read();
252+
if (insertResult != null && !insertResult) {
253+
return;
254+
}
249255

250256
c.output(
251257
String.format(

tests/e2e/redis/basic-ingest-redis-serving.py

Lines changed: 26 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -558,32 +558,6 @@ def try_get_features2():
558558
)
559559

560560

561-
@pytest.mark.timeout(300)
562-
@pytest.mark.run(order=19)
563-
def test_basic_ingest_jobs(client):
564-
# list ingestion jobs given featureset
565-
cust_trans_fs = client.get_feature_set(name="customer_transactions")
566-
ingest_jobs = client.list_ingest_jobs(
567-
feature_set_ref=FeatureSetRef.from_feature_set(cust_trans_fs)
568-
)
569-
# filter ingestion jobs to only those that are running
570-
ingest_jobs = [
571-
job for job in ingest_jobs if job.status == IngestionJobStatus.RUNNING
572-
]
573-
assert len(ingest_jobs) >= 1
574-
575-
for ingest_job in ingest_jobs:
576-
# restart ingestion ingest_job
577-
client.restart_ingest_job(ingest_job)
578-
ingest_job.wait(IngestionJobStatus.RUNNING)
579-
assert ingest_job.status == IngestionJobStatus.RUNNING
580-
581-
# stop ingestion ingest_job
582-
client.stop_ingest_job(ingest_job)
583-
ingest_job.wait(IngestionJobStatus.ABORTED)
584-
assert ingest_job.status == IngestionJobStatus.ABORTED
585-
586-
587561
@pytest.fixture(scope="module")
588562
def all_types_dataframe():
589563
return pd.DataFrame(
@@ -762,16 +736,33 @@ def test_all_types_ingest_jobs(client, all_types_dataframe):
762736
]
763737
assert len(ingest_jobs) >= 1
764738

765-
for ingest_job in ingest_jobs:
766-
# restart ingestion ingest_job
767-
client.restart_ingest_job(ingest_job)
768-
ingest_job.wait(IngestionJobStatus.RUNNING)
769-
assert ingest_job.status == IngestionJobStatus.RUNNING
739+
ingest_job = ingest_jobs[0]
740+
# restart ingestion ingest_job
741+
# restart means stop current job
742+
# (replacement will be automatically spawned)
743+
client.restart_ingest_job(ingest_job)
744+
# wait for replacement to be created
745+
time.sleep(15) # should be more than polling_interval
746+
747+
# id without timestamp part
748+
# that remains the same between jobs
749+
shared_id = "-".join(ingest_job.id.split("-")[:-1])
750+
replacement_jobs = [
751+
job
752+
for job in ingest_jobs
753+
if job.status == IngestionJobStatus.RUNNING and job.id.startswith(shared_id)
754+
]
770755

771-
# stop ingestion ingest_job
772-
client.stop_ingest_job(ingest_job)
773-
ingest_job.wait(IngestionJobStatus.ABORTED)
774-
assert ingest_job.status == IngestionJobStatus.ABORTED
756+
assert len(replacement_jobs) >= 1
757+
replacement_job = replacement_jobs[0]
758+
759+
replacement_job.wait(IngestionJobStatus.RUNNING)
760+
assert replacement_job.status == IngestionJobStatus.RUNNING
761+
762+
# stop ingestion ingest_job
763+
client.stop_ingest_job(replacement_job)
764+
replacement_job.wait(IngestionJobStatus.ABORTED)
765+
assert replacement_job.status == IngestionJobStatus.ABORTED
775766

776767

777768
@pytest.fixture(scope="module")

0 commit comments

Comments
 (0)