diff --git a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/DeltaSharingCapabilities.java b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/DeltaSharingCapabilities.java index 282007d5..5e6ee61a 100644 --- a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/DeltaSharingCapabilities.java +++ b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/DeltaSharingCapabilities.java @@ -45,9 +45,19 @@ public static ResponseFormat choose(String header) { /** Value for the response {@code delta-sharing-capabilities} header: the format actually used. */ public static String responded(String requestHeader) { - return choose(requestHeader) == ResponseFormat.DELTA - ? "responseformat=delta" - : "responseformat=parquet"; + return responded(requestHeader, false); + } + + public static String responded(String requestHeader, boolean includeEndStreamAction) { + String format = + choose(requestHeader) == ResponseFormat.DELTA + ? "responseformat=delta" + : "responseformat=parquet"; + return includeEndStreamAction ? format + ";includeendstreamaction=true" : format; + } + + public static boolean includeEndStreamAction(String header) { + return "true".equals(parse(header).get("includeendstreamaction")); } /** diff --git a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/ProtocolTables.java b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/ProtocolTables.java index 565e6db0..5f97da6d 100644 --- a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/ProtocolTables.java +++ b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/ProtocolTables.java @@ -172,7 +172,8 @@ public TemporaryTableCredentialsResponse temporaryCredentials( /** * One snapshot of a table in the share that supports url access, as of the latest version, * {@code version}, or an ISO-8601 {@code timestamp} in {@code request}, with presigned file - * URLs. Pushdown hints and startingVersion queries are not implemented. + * URLs. A latest-snapshot query can ask for a refresh token, and a later query can pass it back + * to re-sign the same version. Pushdown hints and startingVersion queries are not implemented. */ public DeltaTableQueryReader.Result query( RecipientPrincipal recipient, @@ -195,12 +196,24 @@ public DeltaTableQueryReader.Result query( if (hasPushdownHint(body)) { throw ApiException.notImplemented("predicate and limit pushdown is not supported"); } + boolean historical = body.version() != null || body.timestamp() != null; + boolean includeRefreshToken = Boolean.TRUE.equals(body.includeRefreshToken()); + if (includeRefreshToken && historical) { + throw ApiException.invalidParameter( + "includeRefreshToken cannot be used when querying a specific version."); + } + if (body.refreshToken() != null && historical) { + throw ApiException.invalidParameter( + "refreshToken cannot be used when querying a specific version."); + } return queryReader.read( resolved, owner(granted), options, body.version(), - parseTimestamp("timestamp", body.timestamp())); + parseTimestamp("timestamp", body.timestamp()), + body.refreshToken(), + includeRefreshToken); } private static boolean hasPushdownHint(QueryTableRequest body) { diff --git a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/QueryTableRequest.java b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/QueryTableRequest.java index b3eaa92c..034b3190 100644 --- a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/QueryTableRequest.java +++ b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/QueryTableRequest.java @@ -15,8 +15,14 @@ public record QueryTableRequest( Long startingVersion, Long endingVersion, Boolean includeHistoricalProtocol, - String idempotencyKey) { + String idempotencyKey, + Boolean includeRefreshToken, + String refreshToken) { static final QueryTableRequest EMPTY = - new QueryTableRequest(null, null, null, null, null, null, null, null, null); + new QueryTableRequest(null, null, null, null, null, null, null, null, null, null, null); + + public QueryTableRequest { + refreshToken = refreshToken == null || refreshToken.isBlank() ? null : refreshToken; + } } diff --git a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/RefreshTokens.java b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/RefreshTokens.java new file mode 100644 index 00000000..82e9f0f4 --- /dev/null +++ b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/RefreshTokens.java @@ -0,0 +1,60 @@ +package io.opensharing.asset.table; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.opensharing.http.ApiException; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.util.Base64; + +/** Encodes snapshot refresh tokens so a later query can resign URLs for the same version. */ +public final class RefreshTokens { + + private static final String PREFIX = "os1rt:"; + private static final ObjectMapper JSON = new ObjectMapper(); + + private RefreshTokens() {} + + public record RefreshToken(String id, long version, long expirationTimestamp) {} + + public static String encode(String tableId, long version, Instant expiration) { + return encode(new RefreshToken(tableId, version, expiration.toEpochMilli())); + } + + public static String encode(RefreshToken token) { + try { + return Base64.getUrlEncoder() + .withoutPadding() + .encodeToString( + (PREFIX + JSON.writeValueAsString(token)).getBytes(StandardCharsets.UTF_8)); + } catch (JsonProcessingException e) { + throw new IllegalStateException("failed to encode refresh token", e); + } + } + + public static RefreshToken decode(String token) { + try { + String decoded = + new String(Base64.getUrlDecoder().decode(token.trim()), StandardCharsets.UTF_8); + if (!decoded.startsWith(PREFIX)) { + throw new IllegalArgumentException("unexpected refresh token"); + } + return JSON.readValue(decoded.substring(PREFIX.length()), RefreshToken.class); + } catch (Exception invalid) { + throw ApiException.invalidParameter("Error decoding refresh token: " + token + "."); + } + } + + public static long versionOf(String token, String tableId) { + RefreshToken decoded = decode(token); + if (decoded.expirationTimestamp() < Instant.now().toEpochMilli()) { + throw ApiException.invalidParameter( + "The refresh token has expired. Please restart the query."); + } + if (!tableId.equals(decoded.id())) { + throw ApiException.invalidParameter( + "The table specified in the refresh token does not match the table being queried."); + } + return decoded.version(); + } +} diff --git a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/TableActions.java b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/TableActions.java index 9c168708..902498ba 100644 --- a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/TableActions.java +++ b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/TableActions.java @@ -117,6 +117,10 @@ public static String deltaFile( new DeltaSingleAction(add)))); } + public static String endStreamAction(String refreshToken, Long minUrlExpirationTimestamp) { + return line(new EndStreamLine(new EndStreamAction(refreshToken, minUrlExpirationTimestamp))); + } + /** * Replace the add path with a signed file URL. When {@code deletionVectorUrl} is set, rewrite the * on-disk deletion vector as a path-type descriptor so clients fetch the signed object. @@ -255,4 +259,9 @@ public record DeltaFile( public record DeltaSingleAction( @JsonSerialize(using = KernelActionSerializer.class) AddFile add) {} + + public record EndStreamLine(EndStreamAction endStreamAction) {} + + @JsonInclude(JsonInclude.Include.NON_NULL) + public record EndStreamAction(String refreshToken, Long minUrlExpirationTimestamp) {} } diff --git a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/delta/DeltaTableQueryReader.java b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/delta/DeltaTableQueryReader.java index 872a5c96..872c2408 100644 --- a/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/delta/DeltaTableQueryReader.java +++ b/server/opensharing-server-core/src/main/java/io/opensharing/asset/table/delta/DeltaTableQueryReader.java @@ -16,6 +16,7 @@ import io.opensharing.asset.table.DeltaSharingCapabilities; import io.opensharing.asset.table.DeltaSharingCapabilities.ResponseFormat; import io.opensharing.asset.table.FileIdHash; +import io.opensharing.asset.table.RefreshTokens; import io.opensharing.asset.table.TableActions; import io.opensharing.asset.table.signer.SignedUrl; import io.opensharing.asset.table.signer.UrlSigners; @@ -31,7 +32,7 @@ /** * Reads Query Table snapshot files from a Delta log. The parquet or delta response is NDJSON: - * protocol, metaData, then one file action per data file. + * protocol, metaData, then one file action per data file, and an endStreamAction when requested. * *

Every file URL is presigned for an hour, or until the catalog's credentials expire if * sooner, so the client reads the data files directly from storage without the provider's @@ -55,18 +56,26 @@ public DeltaTableQueryReader(DeltaKernel kernel, UrlSigners signers) { * @param auth the share owner the catalog vends storage credentials for * @param version the table version to read; null with a null {@code timestamp} reads the latest * @param timestamp read the latest version committed at or before this time + * @param refreshToken a refresh token from an earlier response; reads the version it pins + * @param includeRefreshToken put a refresh token for this snapshot in the endStreamAction */ public Result read( ResolvedAsset table, AuthContext auth, ResponseOptions options, Long version, - Instant timestamp) { + Instant timestamp, + String refreshToken, + boolean includeRefreshToken) { if (version != null && timestamp != null) { throw ApiException.invalidParameter("version and timestamp are mutually exclusive"); } // A time-travel query reports the snapshot version in metaData and on every file. boolean historical = version != null || timestamp != null; + // A refresh token pins the version it was issued for without making the query historical. + if (refreshToken != null) { + version = RefreshTokens.versionOf(refreshToken, table.fullName()); + } DeltaKernel.Session session = kernel.open(table, auth, URL_TTL); Snapshot snapshot = DeltaSnapshots.open(session, version, timestamp); @@ -94,6 +103,14 @@ public Result read( throw new UncheckedIOException(e); } + // A refresh token can only travel in an endStreamAction, so asking for one adds the action. + if (includeRefreshToken || options.includeEndStreamAction()) { + out.endStreamAction( + includeRefreshToken + ? RefreshTokens.encode( + table.fullName(), snapshot.getVersion(), Instant.now().plus(URL_TTL)) + : null); + } return new Result(snapshot.getVersion(), out.ndjson()); } @@ -116,11 +133,14 @@ public record Result(long version, String ndjson) {} * * @param capabilities the request's {@code delta-sharing-capabilities} header; picks the format * @param fileIdHash the parsed {@code fileidhash} header, or null to follow the format + * @param includeEndStreamAction end the response with an endStreamAction */ - public record ResponseOptions(String capabilities, String fileIdHash) {} + public record ResponseOptions( + String capabilities, String fileIdHash, boolean includeEndStreamAction) {} /** - * Builds one NDJSON response. File URLs are presigned with the table's storage credentials. + * Builds one NDJSON response. File URLs are presigned with the table's storage credentials, and + * the earliest expiry among them is reported in the endStreamAction. */ private final class ResponseWriter { @@ -130,6 +150,7 @@ private final class ResponseWriter { private final String capabilities; private final String fileIdHash; private final StringBuilder ndjson = new StringBuilder(); + private Long minUrlExpirationTimestamp; ResponseWriter(DeltaKernel.Session session, ResolvedAsset table, ResponseOptions options) { this.format = DeltaSharingCapabilities.choose(options.capabilities()); @@ -192,6 +213,11 @@ void snapshotFile(Row scanFile, Long version, Long timestamp) { } } + /** {@code refreshToken} may be null. */ + void endStreamAction(String refreshToken) { + ndjson.append(TableActions.endStreamAction(refreshToken, minUrlExpirationTimestamp)); + } + String ndjson() { return ndjson.toString(); } @@ -219,6 +245,10 @@ private SignedFile sign( deletionVectorFileId = FileIdHash.hash(dvPath, fileIdHash, format); expirationTimestamp = Math.min(expirationTimestamp, signedDv.expiration().toEpochMilli()); } + minUrlExpirationTimestamp = + minUrlExpirationTimestamp == null + ? expirationTimestamp + : Math.min(minUrlExpirationTimestamp, expirationTimestamp); return new SignedFile( file.url(), FileIdHash.hash(logPath, fileIdHash, format), diff --git a/server/opensharing-server-core/src/main/java/io/opensharing/protocol/DeltaSharingProtocolService.java b/server/opensharing-server-core/src/main/java/io/opensharing/protocol/DeltaSharingProtocolService.java index 138924dc..d26c3ab3 100644 --- a/server/opensharing-server-core/src/main/java/io/opensharing/protocol/DeltaSharingProtocolService.java +++ b/server/opensharing-server-core/src/main/java/io/opensharing/protocol/DeltaSharingProtocolService.java @@ -135,14 +135,19 @@ public HttpResponse queryTable( QueryTableRequest query = body.isBlank() ? null : ApiJson.MAPPER.readValue(body, QueryTableRequest.class); DeltaTableQueryReader.ResponseOptions options = - new DeltaTableQueryReader.ResponseOptions(capabilities, FileIdHash.parse(fileIdHash)); + new DeltaTableQueryReader.ResponseOptions( + capabilities, + FileIdHash.parse(fileIdHash), + DeltaSharingCapabilities.includeEndStreamAction(capabilities)); DeltaTableQueryReader.Result result = tables.query(RecipientTokenDecorator.recipient(ctx), share, schema, table, query, options); ResponseHeadersBuilder headers = ResponseHeaders.builder(HttpStatus.OK) .contentType(NDJSON) .add(DELTA_TABLE_VERSION, Long.toString(result.version())) - .add(DeltaSharingCapabilities.HEADER, DeltaSharingCapabilities.responded(capabilities)); + .add( + DeltaSharingCapabilities.HEADER, + DeltaSharingCapabilities.responded(capabilities, options.includeEndStreamAction())); if (options.fileIdHash() != null) { headers.add(FileIdHash.HEADER, options.fileIdHash()); } diff --git a/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/DeltaSharingCapabilitiesTest.java b/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/DeltaSharingCapabilitiesTest.java index 813f6e68..ed968a64 100644 --- a/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/DeltaSharingCapabilitiesTest.java +++ b/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/DeltaSharingCapabilitiesTest.java @@ -2,7 +2,9 @@ import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import io.delta.kernel.internal.actions.Protocol; import io.opensharing.asset.table.DeltaSharingCapabilities.ResponseFormat; @@ -19,6 +21,12 @@ void defaultsToParquetWithoutAHeader() { assertEquals("responseformat=parquet", DeltaSharingCapabilities.responded(null)); assertEquals( "responseformat=delta", DeltaSharingCapabilities.responded("responseformat=delta")); + assertEquals( + "responseformat=parquet;includeendstreamaction=true", + DeltaSharingCapabilities.responded(null, true)); + assertFalse(DeltaSharingCapabilities.includeEndStreamAction(null)); + assertTrue( + DeltaSharingCapabilities.includeEndStreamAction("includeendstreamaction=true")); } @Test diff --git a/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/RefreshTokensTest.java b/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/RefreshTokensTest.java new file mode 100644 index 00000000..34472afc --- /dev/null +++ b/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/RefreshTokensTest.java @@ -0,0 +1,34 @@ +package io.opensharing.asset.table; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import io.opensharing.asset.table.RefreshTokens.RefreshToken; +import io.opensharing.http.ApiException; +import java.time.Instant; +import org.junit.jupiter.api.Test; + +class RefreshTokensTest { + + @Test + void encodesAndDecodesRefreshToken() { + RefreshToken token = new RefreshToken("main.sales.orders", 7, 1000); + assertEquals(token, RefreshTokens.decode(RefreshTokens.encode(token))); + } + + @Test + void roundTripsTableIdAndVersion() { + String token = RefreshTokens.encode("main.sales.orders", 7, Instant.now().plusSeconds(60)); + assertEquals(7, RefreshTokens.versionOf(token, "main.sales.orders")); + } + + @Test + void rejectsExpiredAndMismatchedTokens() { + String expired = RefreshTokens.encode("main.sales.orders", 1, Instant.now().minusSeconds(1)); + assertThrows(ApiException.class, () -> RefreshTokens.versionOf(expired, "main.sales.orders")); + String token = RefreshTokens.encode("main.sales.orders", 1, Instant.now().plusSeconds(60)); + assertThrows(ApiException.class, () -> RefreshTokens.versionOf(token, "main.hr.salaries")); + assertThrows( + ApiException.class, () -> RefreshTokens.versionOf("not-a-token", "main.sales.orders")); + } +} diff --git a/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/TableActionsTest.java b/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/TableActionsTest.java index c8f6473c..8a9180ed 100644 --- a/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/TableActionsTest.java +++ b/server/opensharing-server-core/src/test/java/io/opensharing/asset/table/TableActionsTest.java @@ -119,6 +119,21 @@ void encodesParquetAndDeltaFiles() { TableActions.deltaFile("file-1", "dv-1", 1L, null, null, signed)); } + @Test + void encodesEndStreamAction() { + assertEquals( + """ + {"endStreamAction":{"refreshToken":"refresh-token",\ + "minUrlExpirationTimestamp":1652140800000}} + """, + TableActions.endStreamAction("refresh-token", 1652140800000L)); + assertEquals( + """ + {"endStreamAction":{}} + """, + TableActions.endStreamAction(null, null)); + } + @Test void encodesDeltaProtocolAndMetadataFromKernel() { Protocol protocol = protocol(); diff --git a/server/opensharing-server/src/test/java/io/opensharing/protocol/ProtocolTableRefreshIntegrationTest.java b/server/opensharing-server/src/test/java/io/opensharing/protocol/ProtocolTableRefreshIntegrationTest.java new file mode 100644 index 00000000..d0403627 --- /dev/null +++ b/server/opensharing-server/src/test/java/io/opensharing/protocol/ProtocolTableRefreshIntegrationTest.java @@ -0,0 +1,115 @@ +package io.opensharing.protocol; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.linecorp.armeria.common.HttpMethod; +import io.opensharing.TestServer.Response; +import java.util.Arrays; +import java.util.Properties; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; + +/** + * Query Table refresh tokens and endStreamAction on real cloud Delta logs, with presigned file + * URLs. Individual cloud cases are skipped unless their repository credentials are present. + */ +@Timeout(60) +class ProtocolTableRefreshIntegrationTest extends ProtocolApiSupport { + + private static final ObjectMapper JSON = new ObjectMapper(); + + @BeforeAll + static void startCloudServer() { + Properties properties = new Properties(); + properties.setProperty("catalog.local.file", "classpath:local-catalog-cloud.yml"); + restartServer(properties); + } + + @Test + @EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".+") + @EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".+") + void refreshesTable1Query() { + refreshesQuery("table1-refresh", "main.sales.table1", "sales.table1", "2"); + } + + @Test + @EnabledIfEnvironmentVariable(named = "AZURE_TEST_ACCOUNT_KEY", matches = ".+") + void refreshesAzureTableQuery() { + refreshesQuery("azure-refresh", "main.sales.azure", "sales.azure", "0"); + } + + @Test + @EnabledIfEnvironmentVariable(named = "GOOGLE_APPLICATION_CREDENTIALS", matches = ".+") + void refreshesGoogleTableQuery() { + refreshesQuery("gcs-refresh", "main.sales.gcs", "sales.gcs", "0"); + } + + private static void refreshesQuery( + String share, String catalogName, String sharedAs, String latest) { + createShare(share); + addObject(share, "TABLE", catalogName, sharedAs); + String bearer = createAndActivateRecipient(share + "-partner"); + grant(share, share + "-partner"); + String[] alias = sharedAs.split("\\."); + String endpoint = + PROTOCOL + "/shares/" + share + "/schemas/" + alias[0] + "/tables/" + alias[1] + "/query"; + + // includeRefreshToken ends the response with an endStreamAction that holds the token. + Response first = query(endpoint, bearer, "{\"includeRefreshToken\":true}", null); + assertVersion(latest, first); + assertEquals(1, count(first.text(), "endStreamAction"), first.text()); + String[] lines = first.text().split("\n"); + JsonNode endStream = readTree(lines[lines.length - 1]).path("endStreamAction"); + assertTrue(endStream.path("minUrlExpirationTimestamp").isNumber(), first.text()); + String token = endStream.path("refreshToken").asText(); + + // The token re-reads the snapshot it was issued for. + Response refreshed = query(endpoint, bearer, "{\"refreshToken\":\"" + token + "\"}", null); + assertVersion(latest, refreshed); + assertEquals(count(first.text(), "file"), count(refreshed.text(), "file"), refreshed.text()); + + // includeendstreamaction adds an endStreamAction without a token and is echoed back. + Response withEndStream = + query(endpoint, bearer, "{}", "responseformat=delta;includeendstreamaction=true"); + assertVersion(latest, withEndStream); + assertEquals( + "responseformat=delta;includeendstreamaction=true", + withEndStream.headers().get("delta-sharing-capabilities")); + assertEquals(1, count(withEndStream.text(), "endStreamAction"), withEndStream.text()); + assertFalse(withEndStream.text().contains("\"refreshToken\""), withEndStream.text()); + } + + private static JsonNode readTree(String line) { + try { + return JSON.readTree(line); + } catch (JsonProcessingException e) { + throw new AssertionError("not an NDJSON line: " + line, e); + } + } + + /** Lines whose single key is {@code action}, such as file or endStreamAction. */ + private static long count(String ndjson, String action) { + return Arrays.stream(ndjson.split("\n")) + .filter(line -> line.startsWith("{\"" + action + "\"")) + .count(); + } + + private static Response query(String path, String bearer, String body, String capabilities) { + return capabilities == null + ? server.post(path, bearer, body) + : server.call( + HttpMethod.POST, path, bearer, body, "delta-sharing-capabilities", capabilities); + } + + private static void assertVersion(String expected, Response response) { + assertEquals(200, response.status(), response.text()); + assertEquals(expected, response.headers().get("Delta-Table-Version")); + } +} diff --git a/server/opensharing-server/src/test/java/io/opensharing/protocol/ProtocolTableRefreshTest.java b/server/opensharing-server/src/test/java/io/opensharing/protocol/ProtocolTableRefreshTest.java new file mode 100644 index 00000000..66609f3a --- /dev/null +++ b/server/opensharing-server/src/test/java/io/opensharing/protocol/ProtocolTableRefreshTest.java @@ -0,0 +1,96 @@ +package io.opensharing.protocol; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.linecorp.armeria.common.HttpMethod; +import io.opensharing.TestServer.Response; +import io.opensharing.http.ErrorCodes; +import java.util.ArrayList; +import java.util.List; +import org.junit.jupiter.api.Test; + +/** Query Table refresh tokens and endStreamAction, on a local Delta log with no data files. */ +class ProtocolTableRefreshTest extends LocalDeltaTableSupport { + + private static final ObjectMapper JSON = new ObjectMapper(); + + @Test + void refreshesLatestSnapshotQueries() { + createShare("table-refresh"); + addObject("table-refresh", "TABLE", "main.ops.events", "ops.events"); + addObject("table-refresh", "TABLE", "main.sales.orders", "sales.orders"); + String bearer = createAndActivateRecipient("table-refresh-partner"); + grant("table-refresh", "table-refresh-partner"); + String events = PROTOCOL + "/shares/table-refresh/schemas/ops/tables/events/query"; + + // includeRefreshToken ends the response with an endStreamAction that holds the token. + Response first = query(events, bearer, "{\"includeRefreshToken\":true}", null); + assertEquals(200, first.status()); + assertEquals("responseformat=parquet", first.headers().get("delta-sharing-capabilities")); + List lines = lines(first); + assertEquals(3, lines.size()); + String token = lines.get(2).path("endStreamAction").path("refreshToken").asText(); + assertFalse(token.isEmpty()); + + // The token re-reads the version it was issued for, as a latest-snapshot query. + Response refreshed = query(events, bearer, refreshBody(token), null); + assertEquals(200, refreshed.status()); + assertEquals("1", refreshed.headers().get("Delta-Table-Version")); + lines = lines(refreshed); + assertEquals(2, lines.size()); + assertTrue(lines.get(1).path("metaData").path("version").isMissingNode()); + + // A token cannot be combined with a version, nor requested for one. + String pinned = "{\"version\":0,\"refreshToken\":\"" + token + "\"}"; + assertInvalid(query(events, bearer, pinned, null)); + assertInvalid(query(events, bearer, "{\"version\":0,\"includeRefreshToken\":true}", null)); + // A token that does not decode, or was issued for another table, is rejected. + assertInvalid(query(events, bearer, refreshBody("not-a-token"), null)); + String orders = PROTOCOL + "/shares/table-refresh/schemas/sales/tables/orders/query"; + assertInvalid(query(orders, bearer, refreshBody(token), null)); + + // includeendstreamaction adds an endStreamAction without a token and is echoed back. + Response withEndStream = + query(events, bearer, "{}", "responseformat=delta;includeendstreamaction=true"); + assertEquals(200, withEndStream.status()); + assertEquals( + "responseformat=delta;includeendstreamaction=true", + withEndStream.headers().get("delta-sharing-capabilities")); + JsonNode endStream = lines(withEndStream).get(2).path("endStreamAction"); + assertTrue(endStream.isObject()); + assertTrue(endStream.path("refreshToken").isMissingNode()); + } + + private static String refreshBody(String token) { + return "{\"refreshToken\":\"" + token + "\"}"; + } + + private static Response query(String path, String bearer, String body, String capabilities) { + return capabilities == null + ? server.post(path, bearer, body) + : server.call( + HttpMethod.POST, path, bearer, body, "delta-sharing-capabilities", capabilities); + } + + private static List lines(Response response) { + List lines = new ArrayList<>(); + for (String line : response.text().split("\n")) { + try { + lines.add(JSON.readTree(line)); + } catch (JsonProcessingException e) { + throw new AssertionError("not an NDJSON line: " + line, e); + } + } + return lines; + } + + private static void assertInvalid(Response response) { + assertEquals(400, response.status(), response.text()); + assertEquals(ErrorCodes.INVALID_PARAMETER_VALUE, response.errorCode()); + } +}