[rpc] version-gate historical partition lookup

fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java

@@ -28,6 +28,7 @@
 import org.apache.fluss.exception.LeaderNotAvailableException;
 import org.apache.fluss.exception.PartitionNotExistException;
 import org.apache.fluss.exception.RetriableException;
+import org.apache.fluss.exception.UnsupportedVersionException;
 import org.apache.fluss.metadata.PhysicalTablePath;
 import org.apache.fluss.metadata.TableBucket;
 import org.apache.fluss.metadata.TablePartition;
@@ -61,6 +62,7 @@
 
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeLookupRequest;
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makePrefixLookupRequest;
+import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.hasHistoricalLookup;
 
 /**
  * This background thread pool lookup operations from {@link #lookupQueue}, and send lookup requests
@@ -342,12 +344,19 @@ private void sendLookupRequestAndHandleResponse(
             Thread.currentThread().interrupt();
             throw new FlussRuntimeException("interrupted:", e);
         }
+        // Determined from the request rather than from the batches, so that the client and the
+        // server decide the lookup kind by the exact same predicate.
+        boolean historicalRequest = hasHistoricalLookup(lookupRequest);
         gateway.lookup(lookupRequest)
                 .thenAccept(
                         lookupResponse -> {
                             try {
                                 handleLookupResponse(
-                                        tableId, destination, lookupResponse, lookupsByBatchKey);
+                                        tableId,
+                                        destination,
+                                        historicalRequest,
+                                        lookupResponse,
+                                        lookupsByBatchKey);
                             } finally {
                                 maxInFlightReuqestsSemaphore.release();
                             }
@@ -403,6 +412,7 @@ private void sendPrefixLookupRequestAndHandleResponse(
     private void handleLookupResponse(
             long tableId,
             int destination,
+            boolean historicalRequest,
             LookupResponse lookupResponse,
             Map<LookupBatchKey, LookupBatch> lookupsByBatchKey) {
         for (PbLookupRespForBucket pbLookupRespForBucket : lookupResponse.getBucketsRespsList()) {
@@ -413,12 +423,18 @@ private void handleLookupResponse(
                                     ? pbLookupRespForBucket.getPartitionId()
                                     : null,
                             pbLookupRespForBucket.getBucketId());
-            LookupBatchKey lookupBatchKey =
-                    new LookupBatchKey(
-                            tableBucket,
-                            pbLookupRespForBucket.hasOriginalPartitionName()
-                                    ? pbLookupRespForBucket.getOriginalPartitionName()
-                                    : null);
+            String originalPartitionName =
+                    pbLookupRespForBucket.hasOriginalPartitionName()
+                            ? pbLookupRespForBucket.getOriginalPartitionName()
+                            : null;
+
+            if (historicalRequest && originalPartitionName == null) {
+                handleUnEchoedHistoricalResponse(
+                        destination, tableBucket, pbLookupRespForBucket, lookupsByBatchKey);
+                continue;
+            }
+
+            LookupBatchKey lookupBatchKey = new LookupBatchKey(tableBucket, originalPartitionName);
             LookupBatch lookupBatch = lookupsByBatchKey.get(lookupBatchKey);
             if (pbLookupRespForBucket.hasErrorCode()) {
                 ApiError error = ApiError.fromErrorMessage(pbLookupRespForBucket);
@@ -440,6 +456,36 @@ private void handleLookupResponse(
         }
     }
 
+    private void handleUnEchoedHistoricalResponse(
+            int destination,
+            TableBucket tableBucket,
+            PbLookupRespForBucket pbLookupRespForBucket,
+            Map<LookupBatchKey, LookupBatch> lookupsByBatchKey) {
+        ApiError error;
+        if (pbLookupRespForBucket.hasErrorCode()) {
+            error = ApiError.fromErrorMessage(pbLookupRespForBucket);
+        } else {
+            error =
+                    ApiError.fromThrowable(
+                            new UnsupportedVersionException(
+                                    "Server "
+                                            + destination
+                                            + " answered a historical partition lookup on "
+                                            + tableBucket
+                                            + " without echoing the original partition name, so it"
+                                            + " does not support historical partition lookup."
+                                            + " Please upgrade the tablet server to a newer"
+                                            + " version."));
+        }
+        List<LookupQuery> lookups = new ArrayList<>();
+        for (LookupBatch lookupBatch : lookupsByBatchKey.values()) {
+            if (lookupBatch.tableBucket().equals(tableBucket)) {
+                lookups.addAll(lookupBatch.lookups());
+            }
+        }
+        handleLookupError(tableBucket, destination, error, lookups, "lookup");
+    }
+
     private void handlePrefixLookupResponse(
             long tableId,
             int destination,

fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupSenderTest.java

@@ -27,6 +27,7 @@
 import org.apache.fluss.exception.NotLeaderOrFollowerException;
 import org.apache.fluss.exception.TableNotExistException;
 import org.apache.fluss.exception.TimeoutException;
+import org.apache.fluss.exception.UnsupportedVersionException;
 import org.apache.fluss.metadata.PhysicalTablePath;
 import org.apache.fluss.metadata.TableBucket;
 import org.apache.fluss.metadata.TableInfo;
@@ -52,8 +53,10 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
@@ -172,6 +175,60 @@ void testHistoricalLookupsBatchDifferentPartitionsForSameBucket() throws Excepti
                 .containsExactly("dt=20200101", "dt=20200102", "dt=20200103");
     }
 
+    /**
+     * A server that does not support historical lookup answers the request as a normal lookup, so
+     * it succeeds without echoing the original partition name and its values belong to another
+     * keyspace. The negotiated version cannot report this, so the un-echoed response is what the
+     * capability is detected from, and it must not surface as an empty result or an NPE.
+     */
+    @Test
+    void testHistoricalLookupFailsWhenPartitionNameNotEchoed() {
+        gateway.setLookupHandler(this::createNonEchoingResponse);
+
+        LookupQuery query =
+                new LookupQuery(
+                        DATA1_TABLE_PATH_PK, TABLE_BUCKET, bytes("key1"), false, "dt=20200101");
+
+        lookupSender.sendLookups(1, LookupType.LOOKUP, Collections.singletonList(query));
+
+        assertThatThrownBy(() -> query.future().get(5, TimeUnit.SECONDS))
+                .isInstanceOf(ExecutionException.class)
+                .hasRootCauseInstanceOf(UnsupportedVersionException.class)
+                .hasMessageContaining("without echoing the original partition name")
+                .hasMessageContaining("does not support historical partition lookup");
+    }
+
+    /**
+     * Routing validation runs before the historical path on the server and reports once per bucket
+     * without the original partition name, so a server that fully supports historical lookup also
+     * produces un-echoed buckets. Those carry an error code, and the server's own error must reach
+     * every batch routed to the bucket rather than being reported as a missing capability.
+     */
+    @Test
+    void testHistoricalLookupErrorWithoutEchoedNameFailsAllBatchesOnBucket() {
+        gateway.setLookupHandler(
+                request ->
+                        createNonEchoingFailedResponse(
+                                request,
+                                new InvalidBucketRoutingException("invalid bucket routing")));
+
+        LookupQuery query1 =
+                new LookupQuery(
+                        DATA1_TABLE_PATH_PK, TABLE_BUCKET, bytes("key1"), false, "dt=20200101");
+        LookupQuery query2 =
+                new LookupQuery(
+                        DATA1_TABLE_PATH_PK, TABLE_BUCKET, bytes("key2"), false, "dt=20200102");
+
+        lookupSender.sendLookups(1, LookupType.LOOKUP, Arrays.asList(query1, query2));
+
+        for (LookupQuery query : Arrays.asList(query1, query2)) {
+            assertThatThrownBy(() -> query.future().get(5, TimeUnit.SECONDS))
+                    .isInstanceOf(ExecutionException.class)
+                    .hasRootCauseInstanceOf(InvalidBucketRoutingException.class);
+        }
+        assertThat(metadataUpdater.getBucketLocation(TABLE_BUCKET)).isEmpty();
+    }
+
     @Test
     void testNormalAndHistoricalLookupsSplitRequests() throws Exception {
         List<LookupRequest> receivedRequests = Collections.synchronizedList(new ArrayList<>());
@@ -715,6 +772,58 @@ private CompletableFuture<LookupResponse> createPartitionNameEchoResponse(
         return CompletableFuture.completedFuture(response);
     }
 
+    /**
+     * A response from a server that does not implement historical lookup: the unknown request field
+     * is skipped, the request is answered as a normal lookup, and the original partition name is
+     * not echoed.
+     */
+    private CompletableFuture<LookupResponse> createNonEchoingResponse(LookupRequest request) {
+        LookupResponse response = new LookupResponse();
+        for (PbLookupReqForBucket bucketRequest : request.getBucketsReqsList()) {
+            PbLookupRespForBucket bucketResponse = response.addBucketsResp();
+            bucketResponse.setBucketId(bucketRequest.getBucketId());
+            if (bucketRequest.hasPartitionId()) {
+                bucketResponse.setPartitionId(bucketRequest.getPartitionId());
+            }
+            for (int i = 0; i < bucketRequest.getKeysCount(); i++) {
+                bucketResponse
+                        .addValue()
+                        .setValues(
+                                responseValue(
+                                        "",
+                                        new String(
+                                                bucketRequest.getKeyAt(i),
+                                                StandardCharsets.UTF_8)));
+            }
+        }
+        return CompletableFuture.completedFuture(response);
+    }
+
+    /**
+     * An error response without the echoed original partition name, as the server's routing
+     * validation produces it. Routing errors are collected once per bucket, so a bucket appears
+     * once even when several original partitions were batched onto it.
+     */
+    private CompletableFuture<LookupResponse> createNonEchoingFailedResponse(
+            LookupRequest request, Exception exception) {
+        LookupResponse response = new LookupResponse();
+        ApiError error = ApiError.fromThrowable(exception);
+        Set<Integer> respondedBuckets = new HashSet<>();
+        for (PbLookupReqForBucket bucketRequest : request.getBucketsReqsList()) {
+            if (!respondedBuckets.add(bucketRequest.getBucketId())) {
+                continue;
+            }
+            PbLookupRespForBucket bucketResponse = response.addBucketsResp();
+            bucketResponse.setBucketId(bucketRequest.getBucketId());
+            if (bucketRequest.hasPartitionId()) {
+                bucketResponse.setPartitionId(bucketRequest.getPartitionId());
+            }
+            bucketResponse.setErrorCode(error.error().code());
+            bucketResponse.setErrorMessage(error.formatErrMsg());
+        }
+        return CompletableFuture.completedFuture(response);
+    }
+
     private static byte[] bytes(String value) {
         return value.getBytes(StandardCharsets.UTF_8);
     }

fluss-rpc/src/main/java/org/apache/fluss/rpc/protocol/ApiKeys.java

@@ -58,7 +58,8 @@ public enum ApiKeys {
     // Version 0: Uses lake's encoder for primary key encoding (legacy behavior).
     // Version 1: Uses CompactedKeyEncoder for primary key encoding when bucket key differs from
     //            primary key, enabling prefix lookup support.
-    LOOKUP(1017, 0, 1, PUBLIC),
+    // Version 2: Supports original_partition_name in requests and responses for historical lookups.
+    LOOKUP(1017, 0, 2, PUBLIC),
 
     NOTIFY_LEADER_AND_ISR(1018, 0, 0, PRIVATE),
     STOP_REPLICA(1019, 0, 0, PRIVATE),

fluss-rpc/src/test/java/org/apache/fluss/rpc/netty/client/ServerConnectionTest.java

@@ -40,6 +40,7 @@
 import org.apache.fluss.rpc.messages.GetTableSchemaRequest;
 import org.apache.fluss.rpc.messages.ListDatabasesRequest;
 import org.apache.fluss.rpc.messages.LookupRequest;
+import org.apache.fluss.rpc.messages.LookupResponse;
 import org.apache.fluss.rpc.messages.PbApiVersion;
 import org.apache.fluss.rpc.messages.PbLookupReqForBucket;
 import org.apache.fluss.rpc.messages.PbTablePath;
@@ -368,6 +369,37 @@ void testRejectHistoricalWritesForOldServer() throws Exception {
         }
     }
 
+    @Test
+    void testHistoricalLookupNotRejectedForServerAdvertisingLookupV1() throws Exception {
+        nettyServer.close();
+        buildNettyServer(new LookupV1GatewayService());
+
+        ServerConnection connection =
+                new ServerConnection(
+                        bootstrap,
+                        serverNode,
+                        TestingClientMetricGroup.newInstance(),
+                        clientAuthenticator,
+                        (con, ignore) -> {});
+        try {
+            assertThat(connection.send(ApiKeys.LOOKUP, lookupRequest(null)).get())
+                    .isInstanceOf(LookupResponse.class);
+            assertThat(connection.send(ApiKeys.LOOKUP, lookupRequest("dt=20260823")).get())
+                    .isInstanceOf(LookupResponse.class);
+        } finally {
+            connection.close().get();
+        }
+    }
+
+    private static LookupRequest lookupRequest(String originalPartitionName) {
+        LookupRequest request = new LookupRequest().setTableId(1L);
+        request.addBucketsReq().setBucketId(0).addKey(new byte[] {1});
+        if (originalPartitionName != null) {
+            request.getBucketsReqAt(0).setOriginalPartitionName(originalPartitionName);
+        }
+        return request;
+    }
+
     private static PutKvRequest putKvRequest(String originalPartitionName) {
         PutKvRequest request = new PutKvRequest().setTableId(1L).setAcks(1).setTimeoutMs(10_000);
         request.addBucketsReq().setBucketId(0).setRecords(new byte[0]);
@@ -443,6 +475,23 @@ public CompletableFuture<ProduceLogResponse> produceLog(ProduceLogRequest reques
         }
     }
 
+    /** A server that supports historical lookup but still advertises {@code LOOKUP} version 1. */
+    private static class LookupV1GatewayService extends TestingTabletGatewayService {
+        @Override
+        public CompletableFuture<ApiVersionsResponse> apiVersions(ApiVersionsRequest request) {
+            return super.apiVersions(request)
+                    .thenApply(
+                            response -> {
+                                for (PbApiVersion apiVersion : response.getApiVersionsList()) {
+                                    if (apiVersion.getApiKey() == ApiKeys.LOOKUP.id) {
+                                        apiVersion.setMaxVersion(1);
+                                    }
+                                }
+                                return response;
+                            });
+        }
+    }
+
     private static class MockMetricRegistry extends NOPMetricRegistry {
 
         Map<String, Metric> registeredMetrics = new HashMap<>();

fluss-rpc/src/test/java/org/apache/fluss/rpc/protocol/ApiKeysTest.java

@@ -48,4 +48,9 @@ void testSnapshotMetadataSupportsLayoutAwareClient() {
     void testAlterTableSupportsBucketCountChange() {
         assertThat(ApiKeys.ALTER_TABLE.highestSupportedVersion).isEqualTo((short) 1);
     }
+
+    @Test
+    void testLookupAdvertisesHistoricalPartitionSupport() {
+        assertThat(ApiKeys.LOOKUP.highestSupportedVersion).isEqualTo((short) 2);
+    }
 }