Merge branch 'apache:master' into json_core_fix

be/src/olap/tablet_reader.cpp

@@ -578,6 +578,8 @@ void TabletReader::_init_conditions_param_except_leafnode_of_andnode(
         for (int id : read_params.topn_filter_source_node_ids) {
             auto& runtime_predicate =
                     read_params.runtime_state->get_query_ctx()->get_runtime_predicate(id);
+            DCHECK(runtime_predicate.inited())
+                    << "runtime predicate not inited, source_node_id=" << id;
             runtime_predicate.set_tablet_schema(_tablet_schema);
         }
     }

be/src/runtime/runtime_predicate.h

@@ -53,7 +53,7 @@ class RuntimePredicate {
 
     void set_tablet_schema(TabletSchemaSPtr tablet_schema) {
         std::unique_lock<std::shared_mutex> wlock(_rwlock);
-        if (_tablet_schema) {
+        if (_tablet_schema || !_inited) {
             return;
         }
         _tablet_schema = tablet_schema;

be/src/vec/sink/writer/vtablet_writer.cpp

@@ -357,18 +357,20 @@ Status VNodeChannel::init(RuntimeState* state) {
     _cur_add_block_request->set_eos(false);
 
     // add block closure
+    // Has to using value to capture _task_exec_ctx because tablet writer may destroyed during callback.
     _send_block_callback = WriteBlockCallback<PTabletWriterAddBlockResult>::create_shared();
-    _send_block_callback->addFailedHandler([this](bool is_last_rpc) {
-        auto ctx_lock = _task_exec_ctx.lock();
+    _send_block_callback->addFailedHandler([&, task_exec_ctx = _task_exec_ctx](bool is_last_rpc) {
+        auto ctx_lock = task_exec_ctx.lock();
         if (ctx_lock == nullptr) {
             return;
         }
         _add_block_failed_callback(is_last_rpc);
     });
 
     _send_block_callback->addSuccessHandler(
-            [this](const PTabletWriterAddBlockResult& result, bool is_last_rpc) {
-                auto ctx_lock = _task_exec_ctx.lock();
+            [&, task_exec_ctx = _task_exec_ctx](const PTabletWriterAddBlockResult& result,
+                                                bool is_last_rpc) {
+                auto ctx_lock = task_exec_ctx.lock();
                 if (ctx_lock == nullptr) {
                     return;
                 }

fe/fe-core/src/main/java/org/apache/doris/httpv2/rest/LoadAction.java

@@ -23,6 +23,7 @@
 import org.apache.doris.common.Config;
 import org.apache.doris.common.DdlException;
 import org.apache.doris.common.LoadException;
+import org.apache.doris.common.UserException;
 import org.apache.doris.httpv2.entity.ResponseEntityBuilder;
 import org.apache.doris.httpv2.entity.RestBaseResult;
 import org.apache.doris.httpv2.exception.UnauthorizedException;
@@ -362,7 +363,11 @@ private TNetworkAddress selectRedirectBackend(boolean groupCommit) throws LoadEx
     // temporarily addressing the users' needs for audit logs.
     // So this function is not widely tested under general scenario
     private boolean checkClusterToken(String token) {
-        return Env.getCurrentEnv().getLoadManager().getTokenManager().checkAuthToken(token);
+        try {
+            return Env.getCurrentEnv().getLoadManager().getTokenManager().checkAuthToken(token);
+        } catch (UserException e) {
+            throw new UnauthorizedException(e.getMessage());
+        }
     }
 
     // NOTE: This function can only be used for AuditlogPlugin stream load for now.

fe/fe-core/src/main/java/org/apache/doris/load/loadv2/TokenManager.java

@@ -63,11 +63,6 @@ private String generateNewToken() {
         return UUID.randomUUID().toString();
     }
 
-    // this method only will be called in master node, since stream load only send message to master.
-    public boolean checkAuthToken(String token) {
-        return tokenQueue.contains(token);
-    }
-
     public String acquireToken() throws UserException {
         if (Env.getCurrentEnv().isMaster() || FeConstants.runningUnitTest) {
             return tokenQueue.peek();
@@ -81,9 +76,8 @@ public String acquireToken() throws UserException {
         }
     }
 
-    public String acquireTokenFromMaster() throws TException {
+    private String acquireTokenFromMaster() throws TException {
         TNetworkAddress thriftAddress = getMasterAddress();
-
         FrontendService.Client client = getClient(thriftAddress);
 
         if (LOG.isDebugEnabled()) {
@@ -108,7 +102,7 @@ public String acquireTokenFromMaster() throws TException {
             } else {
                 TMySqlLoadAcquireTokenResult result = client.acquireToken();
                 if (result.getStatus().getStatusCode() != TStatusCode.OK) {
-                    throw new TException("commit failed.");
+                    throw new TException("acquire token from master failed. " + result.getStatus());
                 }
                 isReturnToPool = true;
                 return result.getToken();
@@ -122,6 +116,57 @@ public String acquireTokenFromMaster() throws TException {
         }
     }
 
+    /**
+     * Check if the token is valid.
+     * If this is not Master FE, will send the request to Master FE.
+     */
+    public boolean checkAuthToken(String token) throws UserException {
+        if (Env.getCurrentEnv().isMaster() || FeConstants.runningUnitTest) {
+            return tokenQueue.contains(token);
+        } else {
+            try {
+                return checkTokenFromMaster(token);
+            } catch (TException e) {
+                LOG.warn("check token error", e);
+                throw new UserException("Check token from master failed", e);
+            }
+        }
+    }
+
+    private boolean checkTokenFromMaster(String token) throws TException {
+        TNetworkAddress thriftAddress = getMasterAddress();
+        FrontendService.Client client = getClient(thriftAddress);
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Send check token to Master {}", thriftAddress);
+        }
+
+        boolean isReturnToPool = false;
+        try {
+            boolean result = client.checkToken(token);
+            isReturnToPool = true;
+            return result;
+        } catch (TTransportException e) {
+            boolean ok = ClientPool.frontendPool.reopen(client, thriftTimeoutMs);
+            if (!ok) {
+                throw e;
+            }
+            if (e.getType() == TTransportException.TIMED_OUT) {
+                throw e;
+            } else {
+                boolean result = client.checkToken(token);
+                isReturnToPool = true;
+                return result;
+            }
+        } finally {
+            if (isReturnToPool) {
+                ClientPool.frontendPool.returnObject(thriftAddress, client);
+            } else {
+                ClientPool.frontendPool.invalidateObject(thriftAddress, client);
+            }
+        }
+    }
+
 
     private TNetworkAddress getMasterAddress() throws TException {
         Env.getCurrentEnv().checkReadyOrThrowTException();

fe/fe-core/src/main/java/org/apache/doris/service/FrontendServiceImpl.java

@@ -1062,12 +1062,6 @@ private void checkPasswordAndPrivs(String user, String passwd, String db, List<S
         }
     }
 
-    private void checkToken(String token) throws AuthenticationException {
-        if (!Env.getCurrentEnv().getLoadManager().getTokenManager().checkAuthToken(token)) {
-            throw new AuthenticationException("Un matched cluster token.");
-        }
-    }
-
     private void checkPassword(String user, String passwd, String clientIp)
             throws AuthenticationException {
         final String fullUserName = ClusterNamespace.getNameFromFullName(user);
@@ -1132,7 +1126,9 @@ private TLoadTxnBeginResult loadTxnBeginImpl(TLoadTxnBeginRequest request, Strin
                     request.getTbl(),
                     request.getUserIp(), PrivPredicate.LOAD);
         } else {
-            checkToken(request.getToken());
+            if (!checkToken(request.getToken())) {
+                throw new AuthenticationException("Invalid token: " + request.getToken());
+            }
         }
 
         // check label
@@ -1343,7 +1339,9 @@ private void loadTxnPreCommitImpl(TLoadTxnCommitRequest request) throws UserExce
         if (request.isSetAuthCode()) {
             // CHECKSTYLE IGNORE THIS LINE
         } else if (request.isSetToken()) {
-            checkToken(request.getToken());
+            if (!checkToken(request.getToken())) {
+                throw new AuthenticationException("Invalid token: " + request.getToken());
+            }
         } else {
             // refactoring it
             if (CollectionUtils.isNotEmpty(request.getTbls())) {
@@ -2413,6 +2411,20 @@ public TMySqlLoadAcquireTokenResult acquireToken() throws TException {
         return result;
     }
 
+    @Override
+    public boolean checkToken(String token) {
+        String clientAddr = getClientAddrAsString();
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("receive check token request from client: {}", clientAddr);
+        }
+        try {
+            return Env.getCurrentEnv().getLoadManager().getTokenManager().checkAuthToken(token);
+        } catch (Throwable e) {
+            LOG.warn("catch unknown result.", e);
+            return false;
+        }
+    }
+
     @Override
     public TCheckAuthResult checkAuth(TCheckAuthRequest request) throws TException {
         String clientAddr = getClientAddrAsString();

fe/fe-core/src/main/java/org/apache/doris/statistics/util/StatisticsUtil.java

@@ -139,6 +139,9 @@ public static List<ResultRow> execStatisticQuery(String sql) {
             return Collections.emptyList();
         }
         try (AutoCloseConnectContext r = StatisticsUtil.buildConnectContext()) {
+            if (Config.isCloudMode()) {
+                r.connectContext.getCloudCluster();
+            }
             StmtExecutor stmtExecutor = new StmtExecutor(r.connectContext, sql);
             r.connectContext.setExecutor(stmtExecutor);
             return stmtExecutor.executeInternalQuery();
@@ -447,11 +450,25 @@ public static boolean statsTblAvailable() {
         } catch (Throwable t) {
             return false;
         }
-        for (OlapTable table : statsTbls) {
-            for (Partition partition : table.getPartitions()) {
-                if (partition.getBaseIndex().getTablets().stream()
-                        .anyMatch(t -> t.getNormalReplicaBackendIds().isEmpty())) {
-                    return false;
+        if (Config.isCloudMode()) {
+            try (AutoCloseConnectContext r = buildConnectContext()) {
+                r.connectContext.getCloudCluster();
+                for (OlapTable table : statsTbls) {
+                    for (Partition partition : table.getPartitions()) {
+                        if (partition.getBaseIndex().getTablets().stream()
+                                .anyMatch(t -> t.getNormalReplicaBackendIds().isEmpty())) {
+                            return false;
+                        }
+                    }
+                }
+            }
+        } else {
+            for (OlapTable table : statsTbls) {
+                for (Partition partition : table.getPartitions()) {
+                    if (partition.getBaseIndex().getTablets().stream()
+                            .anyMatch(t -> t.getNormalReplicaBackendIds().isEmpty())) {
+                        return false;
+                    }
                 }
             }
         }

gensrc/thrift/FrontendService.thrift

@@ -1462,6 +1462,8 @@ service FrontendService {
 
     TMySqlLoadAcquireTokenResult acquireToken()
 
+    bool checkToken(1: string token)
+
     TConfirmUnusedRemoteFilesResult confirmUnusedRemoteFiles(1: TConfirmUnusedRemoteFilesRequest request)
 
     TCheckAuthResult checkAuth(1: TCheckAuthRequest request)

regression-test/data/data_model_p0/unique/test_unique_table_auto_inc.out

@@ -116,16 +116,18 @@ Carter	500	9994
 Beata	700	9996
 Nereids	900	9998
 
--- !partial_update_value --
-Bob	100	1
-Alice	200	2
-Tom	300	3
-Test	400	4
-Carter	500	5
-Smith	600	6
-Beata	700	7
-Doris	800	8
-Nereids	900	9
+-- !partial_update_value1 --
+Bob	100
+Alice	200
+Tom	300
+Test	400
+Carter	500
+Smith	600
+Beata	700
+Doris	800
+Nereids	900
+
+-- !partial_update_value2 --
 
 -- !partial_update_value --
 Bob	9990	1