(db/snapshot): add more APIs for Snapshot (#2747)

* (db/snapshot): Add more APIs for Snapshot
Signed-off-by: Npeng.xu <peng.xu@zilliz.com>

* (db/snapshot): Refactor field element related code
Signed-off-by: Npeng.xu <peng.xu@zilliz.com>

* (db/snapshot): fix lint error
Signed-off-by: Npeng.xu <peng.xu@zilliz.com>

* (db/snapshot): remove some prints
Signed-off-by: Npeng.xu <peng.xu@zilliz.com>

core/src/db/SSDBImpl.cpp

@@ -223,14 +223,7 @@ SSDBImpl::DropIndex(const std::string& collection_name, const std::string& field
     snapshot::OperationContext context;
     // SS TODO: no lsn for drop index
     context.lsn = ss->GetCollectionCommit()->GetLsn();
-    auto field_element_id = ss->GetFieldElementId(field_name, field_element_name);
-    if (field_element_id == 0) {
-        std::stringstream emsg;
-        emsg << "Invalid field name: \"" << field_name;
-        emsg << "\" or field element name: \"" << field_element_name << "\"";
-        return Status(SS_INVALID_CONTEX_ERROR, emsg.str());
-    }
-    context.stale_field_element = ss->GetResource<snapshot::FieldElement>(field_element_id);
+    STATUS_CHECK(ss->GetFieldElement(field_name, field_element_name, context.stale_field_element));
     auto op = std::make_shared<snapshot::DropAllIndexOperation>(context, ss);
     STATUS_CHECK(op->Push());
 

core/src/db/snapshot/ResourceOperations.cpp

@@ -305,15 +305,10 @@ SegmentFileOperation::SegmentFileOperation(const SegmentFileContext& sc, ScopedS
 
 Status
 SegmentFileOperation::DoExecute(Store& store) {
-    auto field_element_id = GetStartedSS()->GetFieldElementId(context_.field_name, context_.field_element_name);
-    if (field_element_id == 0) {
-        std::stringstream emsg;
-        emsg << GetRepr() << ". Invalid field name: \"" << context_.field_name;
-        emsg << "\" or field element name: \"" << context_.field_element_name << "\"";
-        return Status(SS_INVALID_CONTEX_ERROR, emsg.str());
-    }
-    resource_ = std::make_shared<SegmentFile>(context_.collection_id, context_.partition_id, context_.segment_id,
-                                              field_element_id);
+    FieldElementPtr fe;
+    STATUS_CHECK(GetStartedSS()->GetFieldElement(context_.field_name, context_.field_element_name, fe));
+    resource_ =
+        std::make_shared<SegmentFile>(context_.collection_id, context_.partition_id, context_.segment_id, fe->GetID());
     AddStep(*resource_, false);
     return Status::OK();
 }

core/src/db/snapshot/Snapshot.cpp

@@ -121,6 +121,40 @@ Snapshot::Snapshot(ID_TYPE id) {
     RefAll();
 }
 
+FieldPtr
+Snapshot::GetField(const std::string& name) const {
+    auto it = field_names_map_.find(name);
+    if (it == field_names_map_.end()) {
+        return nullptr;
+    }
+
+    return GetResource<Field>(it->second);
+}
+
+Status
+Snapshot::GetFieldElement(const std::string& field_name, const std::string& field_element_name,
+                          FieldElementPtr& field_element) const {
+    field_element = nullptr;
+    auto itf = field_element_names_map_.find(field_name);
+    if (itf == field_element_names_map_.end()) {
+        std::stringstream emsg;
+        emsg << "Snapshot::GetFieldElement: Specified field \"" << field_name;
+        emsg << "\" not found";
+        return Status(SS_NOT_FOUND_ERROR, emsg.str());
+    }
+
+    auto itfe = itf->second.find(field_element_name);
+    if (itfe == itf->second.end()) {
+        std::stringstream emsg;
+        emsg << "Snapshot::GetFieldElement: Specified field element \"" << field_element_name;
+        emsg << "\" not found";
+        return Status(SS_NOT_FOUND_ERROR, emsg.str());
+    }
+
+    field_element = GetResource<FieldElement>(itfe->second);
+    return Status::OK();
+}
+
 const std::string
 Snapshot::ToString() const {
     auto to_matrix_string = [](const MappingT& mappings, int line_length, size_t ident = 0) -> std::string {

core/src/db/snapshot/Snapshot.h

@@ -112,6 +112,10 @@ class Snapshot : public ReferenceProxy {
         return latest_schema_commit_id_;
     }
 
+    Status
+    GetFieldElement(const std::string& field_name, const std::string& field_element_name,
+                    FieldElementPtr& field_element) const;
+
     // PXU TODO: add const. Need to change Scopedxxxx::Get
     SegmentCommitPtr
     GetSegmentCommitBySegmentId(ID_TYPE segment_id) const {
@@ -178,6 +182,9 @@ class Snapshot : public ReferenceProxy {
         return it != field_names_map_.end();
     }
 
+    FieldPtr
+    GetField(const std::string& name) const;
+
     [[nodiscard]] bool
     HasFieldElement(const std::string& field_name, const std::string& field_element_name) const {
         auto id = GetFieldElementId(field_name, field_element_name);

core/unittest/ssdb/test_db.cpp

@@ -182,8 +182,6 @@ TEST_F(SSDBTest, IndexTest) {
     ASSERT_TRUE(status.ok());
     auto sf_collector = std::make_shared<SegmentFileCollector>(ss, filter1);
     sf_collector->Iterate();
-    std::cout << "Total " << sf_collector->segment_files_.size() << " of field_element_id ";
-    std::cout << field_element_id << std::endl;
     ASSERT_EQ(new_total, sf_collector->segment_files_.size());
 
     status = db_->DropIndex(c1, sf_context.field_name, sf_context.field_element_name);
@@ -193,8 +191,6 @@ TEST_F(SSDBTest, IndexTest) {
     ASSERT_TRUE(status.ok());
     sf_collector = std::make_shared<SegmentFileCollector>(ss, filter1);
     sf_collector->Iterate();
-    std::cout << "Total " << sf_collector->segment_files_.size() << " of field_element_id ";
-    std::cout << field_element_id << std::endl;
     ASSERT_EQ(0, sf_collector->segment_files_.size());
 
     {

core/unittest/ssdb/test_snapshot.cpp

@@ -560,6 +560,14 @@ TEST_F(SnapshotTest, IndexTest) {
     OperationContext d_a_i_ctx;
     d_a_i_ctx.lsn = next_lsn();
     d_a_i_ctx.stale_field_element = ss->GetResource<FieldElement>(field_element_id);
+
+    FieldElement::Ptr fe;
+    status = ss->GetFieldElement(sf_context.field_name, sf_context.field_element_name,
+            fe);
+
+    ASSERT_TRUE(status.ok());
+    ASSERT_EQ(fe, d_a_i_ctx.stale_field_element);
+
     auto drop_all_index_op = std::make_shared<DropAllIndexOperation>(d_a_i_ctx, ss);
     status = drop_all_index_op->Push();
     ASSERT_TRUE(status.ok());