feat: support truncate table statement (#3542)

4paradigm · Nov 15, 2023 · 5731b2b · 5731b2b
1 parent 825d155
commit 5731b2b
Show file tree

Hide file tree

Showing 30 changed files with 661 additions and 70 deletions.
diff --git a/cases/plan/cmd.yaml b/cases/plan/cmd.yaml
@@ -649,6 +649,22 @@ cases:
           +-cmd_type: drop function
           +-if_exists: true
           +-args: [func1]
+  - id: truncate_stmt
+    desc: truncate
+    sql: TRUNCATE TABLE t1;
+    expect:
+      node_tree_str: |
+        +-node[CMD]
+          +-cmd_type: truncate table
+          +-args: [t1]
+  - id: truncate_stmt_db
+    desc: truncate
+    sql: TRUNCATE TABLE db1.t1;
+    expect:
+      node_tree_str: |
+        +-node[CMD]
+          +-cmd_type: truncate table
+          +-args: [db1, t1]
   - id: exit_stmt
     desc: exit statement
     sql: EXIT;

diff --git a/docs/en/reference/sql/ddl/TRUNCATE_TABLE_STATEMENT.md b/docs/en/reference/sql/ddl/TRUNCATE_TABLE_STATEMENT.md
@@ -0,0 +1,16 @@
+# TRUNCATE TABLE
+
+```
+TRUNCATE TABLE table_name
+```
+
+`TRUNCATE TABLE` statement is used to clear the specified table.
+
+## Example: clear t1
+
+```sql
+TRUNCATE TABLE t1;
+-- Truncate table t1? yes/no
+-- yes
+-- SUCCEED
+```
diff --git a/docs/en/reference/sql/ddl/index.rst b/docs/en/reference/sql/ddl/index.rst
@@ -24,3 +24,4 @@ Data Definition Statement (DDL)
     SHOW_FUNCTIONS
     DROP_FUNCTION
     SHOW_CREATE_TABLE_STATEMENT
+    TRUNCATE_TABLE_STATEMENT
diff --git a/docs/zh/openmldb_sql/ddl/TRUNCATE_TABLE_STATEMENT.md b/docs/zh/openmldb_sql/ddl/TRUNCATE_TABLE_STATEMENT.md
@@ -0,0 +1,16 @@
+# TRUNCATE TABLE
+
+```
+TRUNCATE TABLE table_name
+```
+
+`TRUNCATE TABLE`语句用清空指定的表。
+
+## Example: 清空t1表
+
+```sql
+TRUNCATE TABLE t1;
+-- Truncate table t1? yes/no
+-- yes
+-- SUCCEED
+```
diff --git a/docs/zh/openmldb_sql/ddl/index.rst b/docs/zh/openmldb_sql/ddl/index.rst
@@ -24,3 +24,4 @@
     SHOW_FUNCTIONS
     DROP_FUNCTION
     SHOW_CREATE_TABLE_STATEMENT
+    TRUNCATE_TABLE_STATEMENT
diff --git a/hybridse/include/node/node_enum.h b/hybridse/include/node/node_enum.h
@@ -285,6 +285,7 @@ enum CmdType {
     kCmdDropFunction,
     kCmdShowJobLog,
     kCmdShowCreateTable,
+    kCmdTruncate,
     kCmdFake,  // not a real cmd, for testing purpose only
     kLastCmd = kCmdFake,
 };

diff --git a/hybridse/src/node/sql_node.cc b/hybridse/src/node/sql_node.cc
@@ -76,6 +76,7 @@ static absl::flat_hash_map<CmdType, absl::string_view> CreateCmdTypeNamesMap() {
         {CmdType::kCmdDropFunction, "drop function"},
         {CmdType::kCmdShowFunctions, "show functions"},
         {CmdType::kCmdShowJobLog, "show joblog"},
+        {CmdType::kCmdTruncate, "truncate table"},
     };
     for (auto kind = 0; kind < CmdType::kLastCmd; ++kind) {
         DCHECK(map.find(static_cast<CmdType>(kind)) != map.end());

diff --git a/hybridse/src/planv2/ast_node_converter.cc b/hybridse/src/planv2/ast_node_converter.cc
@@ -611,6 +611,16 @@ base::Status ConvertStatement(const zetasql::ASTStatement* statement, node::Node
             *output = node;
             break;
         }
+        case zetasql::AST_TRUNCATE_STATEMENT: {
+            const zetasql::ASTTruncateStatement* truncate_statement =
+                statement->GetAsOrNull<zetasql::ASTTruncateStatement>();
+            std::vector<std::string> names;
+            CHECK_STATUS(AstPathExpressionToStringList(truncate_statement->target_path(), names));
+            auto node =
+                dynamic_cast<node::CmdNode*>(node_manager->MakeCmdNode(node::CmdType::kCmdTruncate, names));
+            *output = node;
+            break;
+        }
         case zetasql::AST_DROP_FUNCTION_STATEMENT: {
             const zetasql::ASTDropFunctionStatement* drop_fun_statement =
                 statement->GetAsOrNull<zetasql::ASTDropFunctionStatement>();

diff --git a/src/base/status.h b/src/base/status.h
@@ -93,6 +93,7 @@ enum ReturnCode {
     kExceedMaxMemory = 160,
     kInvalidArgs = 161,
     kCheckIndexFailed = 162,
+    kCatalogUpdateFailed = 163,
     kNameserverIsNotLeader = 300,
     kAutoFailoverIsEnabled = 301,
     kEndpointIsNotExist = 302,
@@ -127,7 +128,10 @@ enum ReturnCode {
     kCheckParameterFailed = 331,
     kCreateProcedureFailedOnTablet = 332,
     kCreateFunctionFailedOnTablet = 333,
-    kOPAlreadyExists = 317,
+    kOPAlreadyExists = 334,
+    kOffsetMismatch = 335,
+    kGetTabletFailed = 336,
+    kTruncateTableFailed = 337,
     kReplicaClusterAliasDuplicate = 400,
     kConnectRelicaClusterZkFailed = 401,
     kNotSameReplicaName = 402,

diff --git a/src/catalog/tablet_catalog.cc b/src/catalog/tablet_catalog.cc
@@ -213,7 +213,7 @@ void TabletTableHandler::AddTable(std::shared_ptr<::openmldb::storage::Table> ta
     do {
         old_tables = std::atomic_load_explicit(&tables_, std::memory_order_acquire);
         new_tables = std::make_shared<Tables>(*old_tables);
-        new_tables->emplace(table->GetPid(), table);
+        new_tables->insert_or_assign(table->GetPid(), table);
     } while (!atomic_compare_exchange_weak(&tables_, &old_tables, new_tables));
 }
 

diff --git a/src/client/ns_client.cc b/src/client/ns_client.cc
@@ -297,6 +297,19 @@ bool NsClient::DropTable(const std::string& db, const std::string& name, std::st
     return false;
 }
 
+base::Status NsClient::TruncateTable(const std::string& db, const std::string& name) {
+    ::openmldb::nameserver::TruncateTableRequest request;
+    request.set_name(name);
+    request.set_db(db);
+    ::openmldb::nameserver::TruncateTableResponse response;
+    bool ok = client_.SendRequest(&::openmldb::nameserver::NameServer_Stub::TruncateTable, &request, &response,
+                                  FLAGS_request_timeout_ms, 1);
+    if (ok && response.code() == 0) {
+        return {};
+    }
+    return {response.code(), response.msg()};
+}
+
 bool NsClient::SyncTable(const std::string& name, const std::string& cluster_alias, uint32_t pid, std::string& msg) {
     ::openmldb::nameserver::SyncTableRequest request;
     request.set_name(name);

diff --git a/src/client/ns_client.h b/src/client/ns_client.h
@@ -112,6 +112,8 @@ class NsClient : public Client {
     bool DropTable(const std::string& db, const std::string& name,
                    std::string& msg);  // NOLINT
 
+    base::Status TruncateTable(const std::string& db, const std::string& name);
+
     bool SyncTable(const std::string& name, const std::string& cluster_alias, uint32_t pid,
                    std::string& msg);  // NOLINT
 

diff --git a/src/client/tablet_client.cc b/src/client/tablet_client.cc
@@ -153,6 +153,20 @@ bool TabletClient::SQLBatchRequestQuery(const std::string& db, const std::string
     return true;
 }
 
+base::Status TabletClient::TruncateTable(uint32_t tid, uint32_t pid) {
+    ::openmldb::api::TruncateTableRequest request;
+    ::openmldb::api::TruncateTableResponse response;
+    request.set_tid(tid);
+    request.set_pid(pid);
+    if (!client_.SendRequest(&::openmldb::api::TabletServer_Stub::TruncateTable, &request, &response,
+                                  FLAGS_request_timeout_ms, 1)) {
+        return {base::ReturnCode::kRPCError, "send request failed!"};
+    } else if (response.code() == 0) {
+        return {};
+    }
+    return {response.code(), response.msg()};
+}
+
 base::Status TabletClient::CreateTable(const ::openmldb::api::TableMeta& table_meta) {
     ::openmldb::api::CreateTableRequest request;
     ::openmldb::api::TableMeta* table_meta_ptr = request.mutable_table_meta();

diff --git a/src/client/tablet_client.h b/src/client/tablet_client.h
@@ -56,6 +56,8 @@ class TabletClient : public Client {
 
     base::Status CreateTable(const ::openmldb::api::TableMeta& table_meta);
 
+    base::Status TruncateTable(uint32_t tid, uint32_t pid);
+
     bool UpdateTableMetaForAddField(uint32_t tid, const std::vector<openmldb::common::ColumnDesc>& cols,
                                     const openmldb::common::VersionPair& pair,
                                     std::string& msg);  // NOLINT

diff --git a/src/cmd/sql_cmd_test.cc b/src/cmd/sql_cmd_test.cc
@@ -1078,6 +1078,47 @@ TEST_P(DBSDKTest, DeployWithBias) {
     ASSERT_TRUE(cs->GetNsClient()->DropDatabase(db, msg));
 }
 
+TEST_P(DBSDKTest, Truncate) {
+    auto cli = GetParam();
+    sr = cli->sr;
+    std::string db_name = "test2";
+    std::string table_name = "test1";
+    std::string ddl = "create table test1 (c1 string, c2 int, c3 bigint, INDEX(KEY=c1, ts=c3));";
+    ProcessSQLs(sr, {
+                        "set @@execute_mode = 'online'",
+                        absl::StrCat("create database ", db_name, ";"),
+                        absl::StrCat("use ", db_name, ";"),
+                        ddl,
+                    });
+    hybridse::sdk::Status status;
+    sr->ExecuteSQL(absl::StrCat("truncate table ", table_name, ";"), &status);
+    ASSERT_TRUE(status.IsOK()) << status.ToString();
+    auto res = sr->ExecuteSQL(absl::StrCat("select * from ", table_name, ";"), &status);
+    ASSERT_EQ(res->Size(), 0);
+    for (int i = 0; i < 10; i++) {
+        std::string key = absl::StrCat("key", i);
+        for (int j = 0; j < 10; j++) {
+            uint64_t ts = 1000 + j;
+            sr->ExecuteSQL(absl::StrCat("insert into ", table_name, " values ('", key, "', 11, ", ts, ");"), &status);
+        }
+    }
+
+    res = sr->ExecuteSQL(absl::StrCat("select * from ", table_name, ";"), &status);
+    ASSERT_EQ(res->Size(), 100);
+    sr->ExecuteSQL(absl::StrCat("truncate table ", table_name, ";"), &status);
+    ASSERT_TRUE(status.IsOK()) << status.ToString();
+    res = sr->ExecuteSQL(absl::StrCat("select * from ", table_name, ";"), &status);
+    ASSERT_EQ(res->Size(), 0);
+    sr->ExecuteSQL(absl::StrCat("insert into ", table_name, " values ('aa', 11, 100);"), &status);
+    res = sr->ExecuteSQL(absl::StrCat("select * from ", table_name, ";"), &status);
+    ASSERT_EQ(res->Size(), 1);
+    ProcessSQLs(sr, {
+                        absl::StrCat("use ", db_name, ";"),
+                        absl::StrCat("drop table ", table_name),
+                        absl::StrCat("drop database ", db_name),
+                    });
+}
+
 TEST_P(DBSDKTest, DeletetRange) {
     auto cli = GetParam();
     sr = cli->sr;

diff --git a/src/nameserver/name_server_impl.cc b/src/nameserver/name_server_impl.cc
@@ -3716,6 +3716,69 @@ void NameServerImpl::CreateTable(RpcController* controller, const CreateTableReq
     }
 }
 
+void NameServerImpl::TruncateTable(RpcController* controller, const TruncateTableRequest* request,
+        TruncateTableResponse* response, Closure* done) {
+    brpc::ClosureGuard done_guard(done);
+    const std::string& db = request->db();
+    const std::string& name = request->name();
+    std::shared_ptr<::openmldb::nameserver::TableInfo> table_info;
+    {
+        std::lock_guard<std::mutex> lock(mu_);
+        if (!GetTableInfoUnlock(request->name(), request->db(), &table_info)) {
+            PDLOG(WARNING, "table[%s] does not exist in db [%s]", name.c_str(), db.c_str());
+            response->set_code(::openmldb::base::ReturnCode::kTableIsNotExist);
+            response->set_msg("table does not exist");
+            return;
+        }
+        if (IsExistActiveOp(db, name)) {
+            PDLOG(WARNING, "there is active op. db [%s] name [%s]", db.c_str(), name.c_str());
+            response->set_code(::openmldb::base::ReturnCode::kOPAlreadyExists);
+            response->set_msg("there is active op");
+            return;
+        }
+    }
+    uint32_t tid = table_info->tid();
+    for (const auto& partition : table_info->table_partition()) {
+        uint32_t offset = 0;
+        for (const auto& partition_meta : partition.partition_meta()) {
+            if (partition_meta.offset() != offset) {
+                if (offset == 0) {
+                    offset = partition_meta.offset();
+                } else {
+                    PDLOG(WARNING, "table[%s] partition [%d] offset mismatch", name.c_str(), partition.pid());
+                    response->set_code(::openmldb::base::ReturnCode::kOffsetMismatch);
+                    response->set_msg("partition offset mismatch");
+                    return;
+                }
+            }
+        }
+    }
+    for (const auto& partition : table_info->table_partition()) {
+        uint32_t pid = partition.pid();
+        for (const auto& partition_meta : partition.partition_meta()) {
+            const auto& endpoint = partition_meta.endpoint();
+            auto tablet_ptr = GetTablet(endpoint);
+            if (!tablet_ptr) {
+                PDLOG(WARNING, "endpoint[%s] can not find client", endpoint.c_str());
+                response->set_code(::openmldb::base::ReturnCode::kGetTabletFailed);
+                response->set_msg("fail to get client, endpint " + endpoint);
+                return;
+            }
+            auto status = tablet_ptr->client_->TruncateTable(tid, pid);
+            if (!status.OK()) {
+                PDLOG(WARNING, "truncate failed, tid[%u] pid[%u] endpoint[%s] msg [%s]",
+                        tid, pid, endpoint.c_str(), status.GetMsg().c_str());
+                response->set_code(::openmldb::base::ReturnCode::kTruncateTableFailed);
+                response->set_msg(status.GetMsg());
+                return;
+            }
+        }
+    }
+    PDLOG(INFO, "truncate success, db[%s] name[%s]", db.c_str(), name.c_str());
+    response->set_code(::openmldb::base::ReturnCode::kOk);
+    response->set_msg("ok");
+}
+
 bool NameServerImpl::SaveTableInfo(std::shared_ptr<TableInfo> table_info) {
     std::string table_value;
     table_info->SerializeToString(&table_value);
@@ -10571,5 +10634,26 @@ bool NameServerImpl::IsExistActiveOp(const std::string& db, const std::string& n
     return false;
 }
 
+bool NameServerImpl::IsExistActiveOp(const std::string& db, const std::string& name) {
+    for (const auto& op_list : task_vec_) {
+        if (op_list.empty()) {
+            continue;
+        }
+        for (const auto& op_data : op_list) {
+            if (!db.empty() && op_data->op_info_.db() != db) {
+                continue;
+            }
+            if (!name.empty() && op_data->op_info_.name() != name) {
+                continue;
+            }
+            if (op_data->op_info_.task_status() == api::TaskStatus::kInited ||
+                    op_data->op_info_.task_status() == api::TaskStatus::kDoing) {
+                return true;
+            }
+        }
+    }
+    return false;
+}
+
 }  // namespace nameserver
 }  // namespace openmldb
diff --git a/src/nameserver/name_server_impl.h b/src/nameserver/name_server_impl.h
@@ -165,6 +165,9 @@ class NameServerImpl : public NameServer {
     void DropTable(RpcController* controller, const DropTableRequest* request, GeneralResponse* response,
                    Closure* done);
 
+    void TruncateTable(RpcController* controller, const TruncateTableRequest* request,
+            TruncateTableResponse* response, Closure* done);
+
     void AddTableField(RpcController* controller, const AddTableFieldRequest* request, GeneralResponse* response,
                        Closure* done);
 
@@ -688,6 +691,7 @@ class NameServerImpl : public NameServer {
     bool IsExistDataBase(const std::string& db);
 
     bool IsExistActiveOp(const std::string& db, const std::string& name, api::OPType op_type);
+    bool IsExistActiveOp(const std::string& db, const std::string& name);
 
  private:
     std::mutex mu_;

diff --git a/src/proto/name_server.proto b/src/proto/name_server.proto
@@ -121,6 +121,16 @@ message DropTableRequest {
     optional string db = 4 [default = ""];
 }
 
+message TruncateTableRequest {
+    optional string name = 1;
+    optional string db = 2;
+}
+
+message TruncateTableResponse {
+    optional int32 code = 1;
+    optional string msg = 2;
+}
+
 message LoadTableRequest {
     optional string name = 1;
     optional string endpoint = 2;
@@ -531,6 +541,7 @@ message DeploySQLResponse {
 service NameServer {
     rpc CreateTable(CreateTableRequest) returns (GeneralResponse);
     rpc DropTable(DropTableRequest) returns (GeneralResponse);
+    rpc TruncateTable(TruncateTableRequest) returns (TruncateTableResponse);
     rpc ShowTablet(ShowTabletRequest) returns (ShowTabletResponse);
     rpc ShowTable(ShowTableRequest) returns (ShowTableResponse);
     rpc MakeSnapshotNS(MakeSnapshotNSRequest) returns (GeneralResponse);