add featurestoreId to featurestore search doc

logicalclocks · Jun 14, 2024 · 3b6cae4 · 3b6cae4
1 parent 95ac512
commit 3b6cae4
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 2 deletions.
diff --git a/...java/io/hops/hopsworks/api/opensearch/featurestore/OpenSearchFeaturestoreItemBuilder.java b/...java/io/hops/hopsworks/api/opensearch/featurestore/OpenSearchFeaturestoreItemBuilder.java
@@ -22,9 +22,9 @@
 import io.hops.hopsworks.common.featurestore.xattr.dto.FeatureStoreItem;
 import io.hops.hopsworks.common.featurestore.xattr.dto.FeatureViewXAttrDTO;
 import io.hops.hopsworks.common.featurestore.xattr.dto.FeaturegroupXAttr;
-import io.hops.hopsworks.common.opensearch.OpenSearchFeaturestoreHit;
 import io.hops.hopsworks.common.featurestore.xattr.dto.FeaturestoreXAttrsConstants;
 import io.hops.hopsworks.common.featurestore.xattr.dto.TrainingDatasetXAttrDTO;
+import io.hops.hopsworks.common.opensearch.OpenSearchFeaturestoreHit;
 import io.hops.hopsworks.exceptions.GenericException;
 import io.hops.hopsworks.restutils.RESTCodes;
 
@@ -57,6 +57,7 @@ public OpenSearchFeaturestoreItemDTO.Base fromBaseArtifact(OpenSearchFeaturestor
     item.datasetIId = hit.getDatasetIId();
     item.parentProjectId = hit.getProjectId();
     item.parentProjectName = hit.getProjectName();
+    item.featurestoreId = hit.getFeaturestoreId();
     Object featureStoreXAttr = hit.getXattrs().get(FeaturestoreXAttrsConstants.FEATURESTORE);
     if(featureStoreXAttr != null) {
       switch(hit.getDocType()) {
@@ -77,7 +78,6 @@ public OpenSearchFeaturestoreItemDTO.Base fromBaseArtifact(OpenSearchFeaturestor
 
   private OpenSearchFeaturestoreItemDTO.Base populateFeatureStoreItem(OpenSearchFeaturestoreItemDTO.Base item,
                                                                       FeatureStoreItem src){
-    item.featurestoreId = src.getFeaturestoreId();
     item.description = src.getDescription();
     item.created = new Date(src.getCreateDate());
     item.creator = new UserDTO(userFacade.findByEmail(src.getCreator()));

diff --git a/...common/src/main/java/io/hops/hopsworks/common/commands/featurestore/search/SearchDoc.java b/...common/src/main/java/io/hops/hopsworks/common/commands/featurestore/search/SearchDoc.java
@@ -40,6 +40,9 @@ public class SearchDoc {
   @JsonProperty(FeaturestoreXAttrsConstants.DATASET_INODE_ID)
   @JsonInclude(JsonInclude.Include.NON_NULL)
   private Long datasetIId;
+  @JsonProperty(FeaturestoreXAttrsConstants.FEATURESTORE_ID)
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  private Integer featurestoreId;
   @JsonProperty(FeaturestoreXAttrsConstants.OPENSEARCH_XATTR)
   @JsonInclude(JsonInclude.Include.NON_NULL)
   private XAttr xattr;
@@ -92,6 +95,14 @@ public void setDatasetIId(Long datasetIId) {
     this.datasetIId = datasetIId;
   }
 
+  public Integer getFeaturestoreId() {
+    return featurestoreId;
+  }
+
+  public void setFeaturestoreId(Integer featurestoreId) {
+    this.featurestoreId = featurestoreId;
+  }
+
   public XAttr getXattr() {
     return xattr;
   }

diff --git a/...a/io/hops/hopsworks/common/commands/featurestore/search/SearchFSOpenSearchController.java b/...a/io/hops/hopsworks/common/commands/featurestore/search/SearchFSOpenSearchController.java
@@ -17,6 +17,7 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.hops.hopsworks.common.commands.CommandException;
+import io.hops.hopsworks.common.featurestore.FeaturestoreController;
 import io.hops.hopsworks.common.featurestore.featuregroup.FeaturegroupController;
 import io.hops.hopsworks.common.featurestore.metadata.FeatureStoreKeywordControllerIface;
 import io.hops.hopsworks.common.featurestore.metadata.FeatureStoreTagControllerIface;
@@ -40,9 +41,11 @@
 import io.hops.hopsworks.exceptions.OpenSearchException;
 import io.hops.hopsworks.exceptions.ServiceException;
 import io.hops.hopsworks.persistence.entity.commands.search.SearchFSCommand;
+import io.hops.hopsworks.persistence.entity.featurestore.Featurestore;
 import io.hops.hopsworks.persistence.entity.featurestore.featureview.FeatureView;
 import io.hops.hopsworks.persistence.entity.featurestore.metadata.FeatureStoreTag;
 import io.hops.hopsworks.persistence.entity.featurestore.trainingdataset.TrainingDatasetFeature;
+import io.hops.hopsworks.persistence.entity.project.Project;
 import io.hops.hopsworks.restutils.RESTCodes;
 import org.opensearch.action.delete.DeleteRequest;
 import org.opensearch.action.index.IndexRequest;
@@ -79,6 +82,8 @@ public class SearchFSOpenSearchController {
   @EJB
   private FeaturegroupController featureGroupCtrl;
   @EJB
+  private FeaturestoreController featurestoreController;
+  @EJB
   private TrainingDatasetController trainingDatasetCtrl;
   @EJB
   private InodeController inodeCtrl;
@@ -158,6 +163,8 @@ private SearchDoc create(SearchFSCommand c) throws CommandException {
     String featureStorePath = Utils.getFeaturestorePath(c.getProject(), settings);
     Long featureStoreInode = inodeCtrl.getInodeAtPath(featureStorePath).getId();
     doc.setDatasetIId(featureStoreInode);
+    Featurestore featurestore = getFeatureStoreForProject(c.getProject());
+    doc.setFeaturestoreId(featurestore.getId());
     if(c.getFeatureGroup() != null) {
       doc.setDocType(OpenSearchDocType.FEATURE_GROUP);
       doc.setName(c.getFeatureGroup().getName());
@@ -176,6 +183,15 @@ private SearchDoc create(SearchFSCommand c) throws CommandException {
     return doc;
   }
 
+  private Featurestore getFeatureStoreForProject(Project project) throws CommandException {
+    try {
+      return featurestoreController.getProjectFeaturestore(project);
+    } catch (FeaturestoreException e) {
+      throw new CommandException(RESTCodes.CommandErrorCode.INTERNAL_SERVER_ERROR, Level.FINE,
+        e.getUsrMsg(), e.getDevMsg());
+    }
+  }
+
   private SearchDoc updateTags(SearchFSCommand c) throws CommandException {
     SearchDoc doc =  new SearchDoc();
     Map<String, FeatureStoreTag> tags;

diff --git a/...s-common/src/main/java/io/hops/hopsworks/common/opensearch/OpenSearchFeaturestoreHit.java b/...s-common/src/main/java/io/hops/hopsworks/common/opensearch/OpenSearchFeaturestoreHit.java
@@ -46,6 +46,7 @@ public class OpenSearchFeaturestoreHit implements Comparator<OpenSearchFeaturest
   private Integer projectId;
   private String projectName;
   private Long datasetIId;
+  private Integer featurestoreId;
   private Map<String, Object> xattrs = new HashMap<>();
 
   public OpenSearchFeaturestoreHit() {
@@ -72,6 +73,8 @@ public static OpenSearchFeaturestoreHit instance(SearchHit hit) throws OpenSearc
           feHit.projectName = entry.getValue().toString();
         } else if (entry.getKey().equals(FeaturestoreXAttrsConstants.DATASET_INODE_ID)) {
           feHit.datasetIId = Long.parseLong(entry.getValue().toString());
+        } else if (entry.getKey().equals(FeaturestoreXAttrsConstants.FEATURESTORE_ID)) {
+          feHit.featurestoreId = Integer.parseInt(entry.getValue().toString());
         } else if (entry.getKey().equals(FeaturestoreXAttrsConstants.OPENSEARCH_XATTR)) {
           feHit.xattrs = (Map)entry.getValue();
         }
@@ -158,6 +161,14 @@ public Map<String, Object> getXattrs() {
     return xattrs;
   }
 
+  public Integer getFeaturestoreId() {
+    return featurestoreId;
+  }
+
+  public void setFeaturestoreId(Integer featurestoreId) {
+    this.featurestoreId = featurestoreId;
+  }
+
   public void setXattrs(Map<String, Object> xattrs) {
     this.xattrs = xattrs;
   }