Skip to content

Commit f91fcba

Browse files
bevzzzbevzzz
committed
feat: add async implementation
1 parent f14289f commit f91fcba

6 files changed

Lines changed: 365 additions & 26 deletions

File tree

‎src/it/java/io/weaviate/integration/BackupITest.java‎

Lines changed: 104 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import java.util.UUID;
77
import java.util.concurrent.CompletableFuture;
88
import java.util.concurrent.CompletionException;
9+
import java.util.concurrent.ExecutionException;
910
import java.util.concurrent.TimeoutException;
1011
import java.util.stream.IntStream;
1112

@@ -48,7 +49,7 @@ public void test_lifecycle() throws IOException, TimeoutException {
4849
// Act: start backup
4950
var started = client.backup.create(backup_1, backend,
5051
backup -> backup
51-
.excludeCollections(nsC, nsBig)
52+
.includeCollections(nsA, nsB)
5253
.compressionLevel(CompressionLevel.BEST_SPEED));
5354

5455
// Assert
@@ -104,9 +105,8 @@ public void test_lifecycle() throws IOException, TimeoutException {
104105
// Assert: all 3 backups are present
105106
var all = client.backup.list(backend);
106107
Assertions.assertThat(all).as("all backups")
107-
.hasSize(3)
108108
.extracting(Backup::id)
109-
.containsOnly(backup_1, backup_2, backup_3);
109+
.contains(backup_1, backup_2, backup_3);
110110

111111
// Act: delete data and restore backup #1
112112
client.collections.delete(nsA);
@@ -120,11 +120,107 @@ public void test_lifecycle() throws IOException, TimeoutException {
120120
Assertions.assertThat(collectionA.size()).as("after restore backup #1").isEqualTo(1);
121121
}
122122

123-
@Test(expected = IllegalStateException.class)
124-
public void test_cancelRestore() throws IOException {
125-
var backup = new Backup("#1", "/tmp/bak/#1", "filesystem", List.of("Things"), BackupStatus.STARTED, null,
126-
Backup.Operation.RESTORE);
127-
backup.cancel(client);
123+
@Test
124+
public void test_lifecycle_async() throws ExecutionException, InterruptedException, Exception {
125+
// Arrange
126+
String nsA = ns("A"), nsB = ns("B"), nsC = ns("C"), nsBig = ns("Big");
127+
String backup_1 = ns("backup_1").toLowerCase();
128+
String backend = "filesystem";
129+
130+
try (final var async = client.async()) {
131+
// Start writing data in the background so it's ready
132+
// by the time we get to backup #3.
133+
var spam = spamData(nsBig);
134+
135+
CompletableFuture.allOf(
136+
async.collections.create(nsA),
137+
async.collections.create(nsB),
138+
async.collections.create(nsC))
139+
.join();
140+
141+
// Insert some data to check restore later
142+
var collectionA = async.collections.use(nsA);
143+
collectionA.data.insert(Map.of()).join();
144+
145+
// Act: start backup
146+
var started = async.backup.create(backup_1, backend,
147+
backup -> backup
148+
.includeCollections(nsA, nsB)
149+
.compressionLevel(CompressionLevel.BEST_SPEED))
150+
.join();
151+
152+
// Assert
153+
Assertions.assertThat(started)
154+
.as("created backup operation")
155+
.returns(backup_1, Backup::id)
156+
.returns(backend, Backup::backend)
157+
.returns(BackupStatus.STARTED, Backup::status)
158+
.returns(null, Backup::error)
159+
.extracting(Backup::includesCollections, InstanceOfAssertFactories.list(String.class))
160+
.containsOnly(nsA, nsB);
161+
162+
// Act: await backup competion
163+
var completed = started.waitForCompletion(async).join();
164+
165+
// Assert
166+
Assertions.assertThat(completed)
167+
.as("await backup completion")
168+
.returns(backup_1, Backup::id)
169+
.returns(backend, Backup::backend)
170+
.returns(BackupStatus.SUCCESS, Backup::status)
171+
.returns(null, Backup::error);
172+
173+
// Act: create another backup
174+
String backup_2 = ns("backup_2").toLowerCase();
175+
async.backup.create(backup_2, backend)
176+
.thenCompose(bak -> bak.waitForCompletion(async))
177+
.join();
178+
179+
// Assert: check the second backup is created successfully
180+
var status_2 = async.backup.getCreateStatus(backup_2, backend).join();
181+
Assertions.assertThat(status_2).as("backup #2").get()
182+
.returns(BackupStatus.SUCCESS, Backup::status);
183+
184+
// Act: create and cancel
185+
// Try to throttle this backup by creating a lot of objects,
186+
// limiting CPU resources and requiring high compression ratio.
187+
// This is to avoid flaky tests and make sure we can cancel
188+
// the backup before it completes successfully.
189+
String backup_3 = ns("backup_3").toLowerCase();
190+
spam.join();
191+
async.backup.create(backup_3, backend,
192+
backup -> backup
193+
.includeCollections(nsA, nsB, nsC, nsBig)
194+
.cpuPercentage(1)
195+
.compressionLevel(CompressionLevel.BEST_COMPRESSION))
196+
.thenCompose(cancelMe -> cancelMe.cancel(async)
197+
.thenCompose(__ -> cancelMe.waitForStatus(async, BackupStatus.CANCELED,
198+
wait -> wait.interval(500))))
199+
.join();
200+
201+
// Assert: check the backup is cancelled
202+
var status_3 = async.backup.getCreateStatus(backup_3, backend).join();
203+
Assertions.assertThat(status_3).as("backup #3").get()
204+
.returns(BackupStatus.CANCELED, Backup::status);
205+
206+
// Assert: all 3 backups are present
207+
var all = async.backup.list(backend).join();
208+
Assertions.assertThat(all).as("all backups")
209+
.extracting(Backup::id)
210+
.contains(backup_1, backup_2, backup_3);
211+
212+
// Act: delete data and restore backup #1
213+
async.collections.delete(nsA).join();
214+
async.backup.restore(backup_1, backend, restore -> restore.includeCollections(nsA)).join();
215+
216+
// Assert: object inserted in the beginning of the test is present
217+
var restore_1 = async.backup.getRestoreStatus(backup_1, backend)
218+
.thenCompose(bak -> bak.orElseThrow().waitForCompletion(async))
219+
.join();
220+
Assertions.assertThat(restore_1).as("restore backup #1")
221+
.returns(BackupStatus.SUCCESS, Backup::status);
222+
Assertions.assertThat(collectionA.size().join()).as("after restore backup #1").isEqualTo(1);
223+
}
128224
}
129225

130226
@Test(expected = IllegalStateException.class)

‎src/main/java/io/weaviate/client6/v1/api/WeaviateClient.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ public class WeaviateClient implements AutoCloseable {
3232
/** Client for {@code /aliases} endpoints for managing collection aliases. */
3333
public final WeaviateAliasClient alias;
3434

35-
/** Client for {@code /backups} endpoints for managing collection aliases. */
35+
/** Client for {@code /backups} endpoints for managing backups. */
3636
public final WeaviateBackupClient backup;
3737

3838
public WeaviateClient(Config config) {

‎src/main/java/io/weaviate/client6/v1/api/WeaviateClientAsync.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import java.util.function.Function;
66

77
import io.weaviate.client6.v1.api.alias.WeaviateAliasClientAsync;
8+
import io.weaviate.client6.v1.api.backup.WeaviateBackupClientAsync;
89
import io.weaviate.client6.v1.api.collections.WeaviateCollectionsClient;
910
import io.weaviate.client6.v1.api.collections.WeaviateCollectionsClientAsync;
1011
import io.weaviate.client6.v1.internal.ObjectBuilder;
@@ -30,6 +31,9 @@ public class WeaviateClientAsync implements AutoCloseable {
3031
/** Client for {@code /aliases} endpoints for managing collection aliases. */
3132
public final WeaviateAliasClientAsync alias;
3233

34+
/** Client for {@code /backups} endpoints for managing backups. */
35+
public final WeaviateBackupClientAsync backup;
36+
3337
/**
3438
* This constructor is blocking if {@link Authentication} configured,
3539
* as the client will need to do the initial token exchange.
@@ -84,6 +88,7 @@ public WeaviateClientAsync(Config config) {
8488
this.restTransport = _restTransport;
8589
this.grpcTransport = new DefaultGrpcTransport(grpcOpt);
8690
this.alias = new WeaviateAliasClientAsync(restTransport);
91+
this.backup = new WeaviateBackupClientAsync(restTransport);
8792
this.collections = new WeaviateCollectionsClientAsync(restTransport, grpcTransport);
8893
}
8994

‎src/main/java/io/weaviate/client6/v1/api/backup/Backup.java‎

Lines changed: 88 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,15 @@
44
import java.util.List;
55
import java.util.Optional;
66
import java.util.concurrent.Callable;
7+
import java.util.concurrent.CompletableFuture;
78
import java.util.concurrent.TimeoutException;
89
import java.util.function.Function;
10+
import java.util.function.Supplier;
911

1012
import com.google.gson.annotations.SerializedName;
1113

1214
import io.weaviate.client6.v1.api.WeaviateClient;
15+
import io.weaviate.client6.v1.api.WeaviateClientAsync;
1316
import io.weaviate.client6.v1.internal.ObjectBuilder;
1417

1518
public record Backup(
@@ -45,8 +48,8 @@ public enum Operation {
4548
/**
4649
* Block until the backup has been created / restored successfully.
4750
*
48-
* @param client Weaviate client. Make sure {@link WeaviateClient#close} is not
49-
* called before this method returns.
51+
* @param client Weaviate client. Make sure {@link WeaviateClient#close}
52+
* is NOT called before this method returns.
5053
* @throws IllegalStateException if {@link #operation} is not set (null).
5154
* @throws TimeoutException in case the wait times out without reaching
5255
* BackupStatus.SUCCESS.
@@ -61,8 +64,8 @@ public Backup waitForCompletion(WeaviateClient client) throws IOException, Timeo
6164
/**
6265
* Block until the backup has been created / restored successfully.
6366
*
64-
* @param client Weaviate client. Make sure {@link WeaviateClient#close} is not
65-
* called before this method returns.
67+
* @param client Weaviate client. Make sure {@link WeaviateClient#close}
68+
* is NOT called before this method returns.
6669
* @param fn Lambda expression for optional parameters.
6770
* @throws IllegalStateException if {@link #operation} is not set (null).
6871
* @throws TimeoutException in case the wait times out without reaching
@@ -79,8 +82,8 @@ public Backup waitForCompletion(WeaviateClient client, Function<WaitOptions.Buil
7982
/**
8083
* Block until the backup operation reaches a certain status.
8184
*
82-
* @param client Weaviate client. Make sure {@link WeaviateClient#close} is not
83-
* called before this method returns.
85+
* @param client Weaviate client. Make sure {@link WeaviateClient#close}
86+
* is NOT called before this method returns.
8487
* @param status Target status.
8588
* @throws IllegalStateException if {@link #operation} is not set (null).
8689
* @throws TimeoutException in case the wait times out without reaching
@@ -96,8 +99,8 @@ public Backup waitForStatus(WeaviateClient client, BackupStatus status) throws I
9699
/**
97100
* Block until the backup operation reaches a certain status.
98101
*
99-
* @param client Weaviate client. Make sure {@link WeaviateClient#close} is not
100-
* called before this method returns.
102+
* @param client Weaviate client. Make sure {@link WeaviateClient#close}
103+
* is NOT called before this method returns.
101104
* @param status Target status.
102105
* @param fn Lambda expression for optional parameters.
103106
* @throws IllegalStateException if {@link #operation} is not set (null).
@@ -117,22 +120,95 @@ public Backup waitForStatus(WeaviateClient client, BackupStatus status,
117120
final Callable<Optional<Backup>> poll = operation == Operation.CREATE
118121
? () -> client.backup.getCreateStatus(id, backend)
119122
: () -> client.backup.getRestoreStatus(id, backend);
120-
return new Waiter(this, poll, options).waitForStatus(status);
123+
return new Waiter(this, options).waitForStatus(status, poll);
121124
}
122125

123126
/**
124127
* Cancel backup creation.
125128
*
126129
* <p>
127-
* This method cannot be called cancel backup restore.
130+
* This method cannot be called to cancel backup restore.
128131
*
129-
* @param client Weaviate client. Make sure {@link WeaviateClient#close} is not
130-
* called before this method returns.
132+
* @param client Weaviate client. Make sure {@link WeaviateClient#close}
133+
* is NOT called before this method returns.
131134
* @throws IOException in case the request was not sent successfully
132135
* due to a malformed request, a networking error
133136
* or the server being unavailable.
134137
*/
135138
public void cancel(WeaviateClient client) throws IOException {
136139
client.backup.cancel(id(), backend());
137140
}
141+
142+
/**
143+
* Poll until backup's been created / restored successfully.
144+
*
145+
* @param client Weaviate client. Make sure {@link WeaviateClientAsync#close}
146+
* is NOT called before this method returns.
147+
* @throws IllegalStateException if {@link #operation} is not set (null).
148+
* @throws TimeoutException in case the wait times out without reaching
149+
* BackupStatus.SUCCESS.
150+
* @throws IOException in case the request was not sent successfully
151+
* due to a malformed request, a networking error
152+
* or the server being unavailable.
153+
*/
154+
public CompletableFuture<Backup> waitForCompletion(WeaviateClientAsync client) {
155+
return waitForStatus(client, BackupStatus.SUCCESS);
156+
}
157+
158+
/**
159+
* Poll until backup's been created / restored successfully.
160+
*
161+
* @param client Weaviate client. Make sure {@link WeaviateClientAsync#close}
162+
* is NOT called before this method returns.
163+
* @param fn Lambda expression for optional parameters.
164+
*/
165+
public CompletableFuture<Backup> waitForCompletion(WeaviateClientAsync client,
166+
Function<WaitOptions.Builder, ObjectBuilder<WaitOptions>> fn) {
167+
return waitForStatus(client, BackupStatus.SUCCESS, fn);
168+
}
169+
170+
/**
171+
* Poll until backup reaches a certain status or the wait times out.
172+
*
173+
* @param client Weaviate client. Make sure {@link WeaviateClientAsync#close}
174+
* is NOT called before this method returns.
175+
* @param status Target status.
176+
*/
177+
public CompletableFuture<Backup> waitForStatus(WeaviateClientAsync client, BackupStatus status) {
178+
return waitForStatus(client, status, ObjectBuilder.identity());
179+
}
180+
181+
/**
182+
* Poll until backup reaches a certain status or the wait times out.
183+
*
184+
* @param client Weaviate client. Make sure {@link WeaviateClientAsync#close}
185+
* is NOT called before this method returns.
186+
* @param status Target status.
187+
* @param fn Lambda expression for optional parameters.
188+
*/
189+
public CompletableFuture<Backup> waitForStatus(WeaviateClientAsync client, BackupStatus status,
190+
Function<WaitOptions.Builder, ObjectBuilder<WaitOptions>> fn) {
191+
if (operation == null) {
192+
throw new IllegalStateException("backup.operation is null");
193+
}
194+
195+
final var options = WaitOptions.of(fn);
196+
final Supplier<CompletableFuture<Optional<Backup>>> poll = operation == Operation.CREATE
197+
? () -> client.backup.getCreateStatus(id, backend)
198+
: () -> client.backup.getRestoreStatus(id, backend);
199+
return new Waiter(this, options).waitForStatusAsync(status, poll);
200+
}
201+
202+
/**
203+
* Cancel backup creation.
204+
*
205+
* <p>
206+
* This method cannot be called to cancel backup restore.
207+
*
208+
* @param client Weaviate client. Make sure {@link WeaviateClientAsync#close}
209+
* is NOT called before this method returns.
210+
*/
211+
public CompletableFuture<Void> cancel(WeaviateClientAsync client) {
212+
return client.backup.cancel(id(), backend());
213+
}
138214
}

0 commit comments

Comments
 (0)