From ecfff5458543ccea4fd022e35a66254d6185c7da Mon Sep 17 00:00:00 2001
From: Jesse Jia <jejia@linkedin.com>
Date: Thu, 10 Oct 2024 18:43:22 -0700
Subject: [PATCH] Add search flag into base clients (#448)

Co-authored-by: Jesse Jia <jejia@jejia-mn2.linkedin.biz>
---
 .../metadata/restli/BaseBrowsableClient.java  | 16 +++++++++
 .../metadata/restli/BaseSearchableClient.java | 34 +++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseBrowsableClient.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseBrowsableClient.java
index 0f44566de..7b05e47c8 100644
--- a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseBrowsableClient.java
+++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseBrowsableClient.java
@@ -37,6 +37,15 @@ public BaseBrowsableClient(@Nonnull Client restliClient) {
   public abstract BrowseResult browse(@Nonnull String inputPath, @Nullable Map<String, String> requestFilters, int from, int size)
       throws RemoteInvocationException;
 
+  /**
+   * Browse method that allows caller to choose the platform.
+   */
+  @Nonnull
+  public BrowseResult browse(@Nonnull String inputPath, @Nullable Map<String, String> requestFilters, @Nullable SearchPlatform searchPlatform,
+      int from, int size) throws RemoteInvocationException {
+    throw new UnsupportedOperationException("Not implemented yet.");
+  }
+
   /**
    * Returns a list of paths for a given urn.
    *
@@ -49,4 +58,11 @@ public StringArray getBrowsePaths(@Nonnull URN urn) throws RemoteInvocationExcep
     throw new UnsupportedOperationException("Not implemented yet.");
   }
 
+  /**
+   * Returns a list of paths for a given urn and allows caller to choose the platform.
+   */
+  @Nonnull
+  public StringArray getBrowsePaths(@Nonnull URN urn, @Nullable SearchPlatform searchPlatform) throws RemoteInvocationException {
+    throw new UnsupportedOperationException("Not implemented yet.");
+  }
 }
\ No newline at end of file
diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableClient.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableClient.java
index 13e4a3d7c..48d39c654 100644
--- a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableClient.java
+++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableClient.java
@@ -19,6 +19,21 @@
  */
 public abstract class BaseSearchableClient<VALUE extends RecordTemplate> extends BaseClient {
 
+  public enum SearchPlatform {
+    ELASTICSEARCH("elasticsearch"),
+    HOSTEDSEARCH("hostedsearch");
+
+    private final String platformName;
+
+    SearchPlatform(String platform) {
+      this.platformName = platform;
+    }
+
+    public String getPlatformName() {
+      return platformName;
+    }
+  }
+
   public BaseSearchableClient(@Nonnull Client restliClient) {
     super(restliClient);
   }
@@ -41,6 +56,16 @@ public abstract CollectionResponse<VALUE> search(@Nonnull String input, @Nullabl
       @Nullable Map<String, String> requestFilters, @Nullable SortCriterion sortCriterion, int start, int count)
       throws RemoteInvocationException;
 
+  /**
+   * Similar to {@link #search(String, StringArray, Map, SortCriterion, int, int)} with platform flag to indicate the request
+   * should be served by elasticsearch or hostedsearch.
+   */
+  public CollectionResponse<VALUE> search(@Nonnull String input, @Nullable StringArray aspectNames,
+      @Nullable Map<String, String> requestFilters, @Nullable SortCriterion sortCriterion, @Nullable SearchPlatform platform,
+      int start, int count) throws RemoteInvocationException {
+    throw new UnsupportedOperationException("Not implemented in BaseSearchableClient.");
+  }
+
   /**
    * Similar to {@link #search(String, StringArray, Map, SortCriterion, int, int)} with null for aspect names, meaning
    * all aspects will be returned.
@@ -51,6 +76,15 @@ public CollectionResponse<VALUE> search(@Nonnull String input, @Nullable Map<Str
     return search(input, null, requestFilters, sortCriterion, start, count);
   }
 
+  /**
+   * Autocomplete with search platform flag.
+   */
+  @Nonnull
+  public AutoCompleteResult autocomplete(@Nonnull String query, @Nullable String field, @Nullable Map<String, String> requestFilters,
+      @Nullable SearchPlatform platform, int limit) throws RemoteInvocationException {
+    throw new UnsupportedOperationException("Not implemented in BaseSearchableClient.");
+  }
+
   /**
    * Autocomplete method that the client will override only if they need this capability.
    *