Datasource disable feature (#2539) (#2552)

Signed-off-by: Vamsi Manohar <[email protected]>
(cherry picked from commit 353b0d7)

core/build.gradle

@@ -79,7 +79,9 @@ jacocoTestCoverageVerification {
             excludes = [
                     'org.opensearch.sql.utils.MLCommonsConstants',
                     'org.opensearch.sql.utils.Constants',
-                    'org.opensearch.sql.datasource.model.*'
+                    'org.opensearch.sql.datasource.model.DataSource',
+                    'org.opensearch.sql.datasource.model.DataSourceStatus',
+                    'org.opensearch.sql.datasource.model.DataSourceType'
             ]
             limit {
                 counter = 'LINE'

core/src/main/java/org/opensearch/sql/datasource/DataSourceService.java

@@ -14,7 +14,8 @@
 public interface DataSourceService {
 
   /**
-   * Returns {@link DataSource} corresponding to the DataSource name.
+   * Returns {@link DataSource} corresponding to the DataSource name only if the datasource is
+   * active and authorized.
    *
    * @param dataSourceName Name of the {@link DataSource}.
    * @return {@link DataSource}.
@@ -40,15 +41,6 @@ public interface DataSourceService {
    */
   DataSourceMetadata getDataSourceMetadata(String name);
 
-  /**
-   * Returns dataSourceMetadata object with specific name. The returned objects contain all the
-   * metadata information without any filtering.
-   *
-   * @param name name of the {@link DataSource}.
-   * @return set of {@link DataSourceMetadata}.
-   */
-  DataSourceMetadata getRawDataSourceMetadata(String name);
-
   /**
    * Register {@link DataSource} defined by {@link DataSourceMetadata}.
    *
@@ -84,4 +76,12 @@ public interface DataSourceService {
    * @param dataSourceName name of the {@link DataSource}.
    */
   Boolean dataSourceExists(String dataSourceName);
+
+  /**
+   * Performs authorization and datasource status check and then returns RawDataSourceMetadata.
+   * Specifically for addressing use cases in SparkQueryDispatcher.
+   *
+   * @param dataSourceName of the {@link DataSource}
+   */
+  DataSourceMetadata verifyDataSourceAccessAndGetRawMetadata(String dataSourceName);
 }

core/src/main/java/org/opensearch/sql/datasource/model/DataSourceMetadata.java

@@ -10,6 +10,7 @@
 import com.fasterxml.jackson.annotation.JsonFormat;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -19,27 +20,26 @@
 import java.util.function.Function;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
-import lombok.Setter;
 import org.apache.commons.lang3.StringUtils;
 import org.opensearch.sql.datasource.DataSourceService;
 
 @Getter
-@Setter
 @EqualsAndHashCode
 @JsonIgnoreProperties(ignoreUnknown = true)
 public class DataSourceMetadata {
 
   public static final String DEFAULT_RESULT_INDEX = "query_execution_result";
   public static final int MAX_RESULT_INDEX_NAME_SIZE = 255;
+  private static String DATASOURCE_NAME_REGEX = "[@*A-Za-z]+?[*a-zA-Z_\\-0-9]*";
   // OS doesn’t allow uppercase: https://tinyurl.com/yse2xdbx
   public static final String RESULT_INDEX_NAME_PATTERN = "[a-z0-9_-]+";
   public static String INVALID_RESULT_INDEX_NAME_SIZE =
       "Result index name size must contains less than "
           + MAX_RESULT_INDEX_NAME_SIZE
-          + " characters";
+          + " characters.";
   public static String INVALID_CHAR_IN_RESULT_INDEX_NAME =
       "Result index name has invalid character. Valid characters are a-z, 0-9, -(hyphen) and"
-          + " _(underscore)";
+          + " _(underscore).";
   public static String INVALID_RESULT_INDEX_PREFIX =
       "Result index must start with " + DEFAULT_RESULT_INDEX;
 
@@ -57,96 +57,188 @@ public class DataSourceMetadata {
 
   @JsonProperty private String resultIndex;
 
+  @JsonProperty private DataSourceStatus status;
+
   public static Function<String, String> DATASOURCE_TO_RESULT_INDEX =
       datasourceName -> String.format("%s_%s", DEFAULT_RESULT_INDEX, datasourceName);
 
-  public DataSourceMetadata(
-      String name,
-      String description,
-      DataSourceType connector,
-      List<String> allowedRoles,
-      Map<String, String> properties,
-      String resultIndex) {
-    this.name = name;
-    String errorMessage = validateCustomResultIndex(resultIndex);
-    if (errorMessage != null) {
-      throw new IllegalArgumentException(errorMessage);
+  private DataSourceMetadata(Builder builder) {
+    this.name = builder.name;
+    this.description = builder.description;
+    this.connector = builder.connector;
+    this.allowedRoles = builder.allowedRoles;
+    this.properties = builder.properties;
+    this.resultIndex = builder.resultIndex;
+    this.status = builder.status;
+  }
+
+  public static class Builder {
+    private String name;
+    private String description;
+    private DataSourceType connector;
+    private List<String> allowedRoles;
+    private Map<String, String> properties;
+    private String resultIndex; // Optional
+    private DataSourceStatus status;
+
+    public Builder() {}
+
+    public Builder(DataSourceMetadata dataSourceMetadata) {
+      this.name = dataSourceMetadata.getName();
+      this.description = dataSourceMetadata.getDescription();
+      this.connector = dataSourceMetadata.getConnector();
+      this.resultIndex = dataSourceMetadata.getResultIndex();
+      this.status = dataSourceMetadata.getStatus();
+      this.allowedRoles = new ArrayList<>(dataSourceMetadata.getAllowedRoles());
+      this.properties = new HashMap<>(dataSourceMetadata.getProperties());
     }
-    if (resultIndex == null) {
-      this.resultIndex = fromNameToCustomResultIndex();
-    } else {
-      this.resultIndex = resultIndex;
+
+    public Builder setName(String name) {
+      this.name = name;
+      return this;
     }
 
-    this.connector = connector;
-    this.description = description;
-    this.properties = properties;
-    this.allowedRoles = allowedRoles;
-  }
+    public Builder setDescription(String description) {
+      this.description = description;
+      return this;
+    }
 
-  public DataSourceMetadata() {
-    this.description = StringUtils.EMPTY;
-    this.allowedRoles = new ArrayList<>();
-    this.properties = new HashMap<>();
-  }
+    public Builder setConnector(DataSourceType connector) {
+      this.connector = connector;
+      return this;
+    }
 
-  /**
-   * Default OpenSearch {@link DataSourceMetadata}. Which is used to register default OpenSearch
-   * {@link DataSource} to {@link DataSourceService}.
-   */
-  public static DataSourceMetadata defaultOpenSearchDataSourceMetadata() {
-    return new DataSourceMetadata(
-        DEFAULT_DATASOURCE_NAME,
-        StringUtils.EMPTY,
-        DataSourceType.OPENSEARCH,
-        Collections.emptyList(),
-        ImmutableMap.of(),
-        null);
-  }
+    public Builder setAllowedRoles(List<String> allowedRoles) {
+      this.allowedRoles = allowedRoles;
+      return this;
+    }
 
-  public String validateCustomResultIndex(String resultIndex) {
-    if (resultIndex == null) {
-      return null;
+    public Builder setProperties(Map<String, String> properties) {
+      this.properties = properties;
+      return this;
     }
-    if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) {
-      return INVALID_RESULT_INDEX_NAME_SIZE;
+
+    public Builder setResultIndex(String resultIndex) {
+      this.resultIndex = resultIndex;
+      return this;
     }
-    if (!resultIndex.matches(RESULT_INDEX_NAME_PATTERN)) {
-      return INVALID_CHAR_IN_RESULT_INDEX_NAME;
+
+    public Builder setDataSourceStatus(DataSourceStatus status) {
+      this.status = status;
+      return this;
     }
-    if (resultIndex != null && !resultIndex.startsWith(DEFAULT_RESULT_INDEX)) {
-      return INVALID_RESULT_INDEX_PREFIX;
+
+    public DataSourceMetadata build() {
+      validateMissingAttributes();
+      validateName();
+      validateCustomResultIndex();
+      fillNullAttributes();
+      return new DataSourceMetadata(this);
     }
-    return null;
-  }
 
-  /**
-   * Since we are using datasource name to create result index, we need to make sure that the final
-   * name is valid
-   *
-   * @param resultIndex result index name
-   * @return valid result index name
-   */
-  private String convertToValidResultIndex(String resultIndex) {
-    // Limit Length
-    if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) {
-      resultIndex = resultIndex.substring(0, MAX_RESULT_INDEX_NAME_SIZE);
+    private void fillNullAttributes() {
+      if (resultIndex == null) {
+        this.resultIndex = fromNameToCustomResultIndex();
+      }
+      if (status == null) {
+        this.status = DataSourceStatus.ACTIVE;
+      }
+      if (description == null) {
+        this.description = StringUtils.EMPTY;
+      }
+      if (properties == null) {
+        this.properties = ImmutableMap.of();
+      }
+      if (allowedRoles == null) {
+        this.allowedRoles = ImmutableList.of();
+      }
     }
 
-    // Pattern Matching: Remove characters that don't match the pattern
-    StringBuilder validChars = new StringBuilder();
-    for (char c : resultIndex.toCharArray()) {
-      if (String.valueOf(c).matches(RESULT_INDEX_NAME_PATTERN)) {
-        validChars.append(c);
+    private void validateMissingAttributes() {
+      List<String> missingAttributes = new ArrayList<>();
+      if (name == null) {
+        missingAttributes.add("name");
+      }
+      if (connector == null) {
+        missingAttributes.add("connector");
+      }
+      if (!missingAttributes.isEmpty()) {
+        String errorMessage =
+            "Datasource configuration error: "
+                + String.join(", ", missingAttributes)
+                + " cannot be null or empty.";
+        throw new IllegalArgumentException(errorMessage);
       }
     }
-    return validChars.toString();
-  }
 
-  public String fromNameToCustomResultIndex() {
-    if (name == null) {
-      throw new IllegalArgumentException("Datasource name cannot be null");
+    private void validateName() {
+      if (!name.matches(DATASOURCE_NAME_REGEX)) {
+        throw new IllegalArgumentException(
+            String.format(
+                "DataSource Name: %s contains illegal characters. Allowed characters:"
+                    + " a-zA-Z0-9_-*@.",
+                name));
+      }
+    }
+
+    private void validateCustomResultIndex() {
+      if (resultIndex == null) {
+        return;
+      }
+      StringBuilder errorMessage = new StringBuilder();
+      if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) {
+        errorMessage.append(INVALID_RESULT_INDEX_NAME_SIZE);
+      }
+      if (!resultIndex.matches(RESULT_INDEX_NAME_PATTERN)) {
+        errorMessage.append(INVALID_CHAR_IN_RESULT_INDEX_NAME);
+      }
+      if (!resultIndex.startsWith(DEFAULT_RESULT_INDEX)) {
+        errorMessage.append(INVALID_RESULT_INDEX_PREFIX);
+      }
+      if (errorMessage.length() > 0) {
+        throw new IllegalArgumentException(errorMessage.toString());
+      }
+    }
+
+    /**
+     * Since we are using datasource name to create result index, we need to make sure that the
+     * final name is valid
+     *
+     * @param resultIndex result index name
+     * @return valid result index name
+     */
+    private String convertToValidResultIndex(String resultIndex) {
+      // Limit Length
+      if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) {
+        resultIndex = resultIndex.substring(0, MAX_RESULT_INDEX_NAME_SIZE);
+      }
+
+      // Pattern Matching: Remove characters that don't match the pattern
+      StringBuilder validChars = new StringBuilder();
+      for (char c : resultIndex.toCharArray()) {
+        if (String.valueOf(c).matches(RESULT_INDEX_NAME_PATTERN)) {
+          validChars.append(c);
+        }
+      }
+      return validChars.toString();
     }
-    return convertToValidResultIndex(DATASOURCE_TO_RESULT_INDEX.apply(name.toLowerCase()));
+
+    private String fromNameToCustomResultIndex() {
+      return convertToValidResultIndex(DATASOURCE_TO_RESULT_INDEX.apply(name.toLowerCase()));
+    }
+  }
+
+  /**
+   * Default OpenSearch {@link DataSourceMetadata}. Which is used to register default OpenSearch
+   * {@link DataSource} to {@link DataSourceService}.
+   */
+  public static DataSourceMetadata defaultOpenSearchDataSourceMetadata() {
+    return new DataSourceMetadata.Builder()
+        .setName(DEFAULT_DATASOURCE_NAME)
+        .setDescription(StringUtils.EMPTY)
+        .setConnector(DataSourceType.OPENSEARCH)
+        .setAllowedRoles(Collections.emptyList())
+        .setProperties(ImmutableMap.of())
+        .build();
   }
 }

core/src/main/java/org/opensearch/sql/datasource/model/DataSourceStatus.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.datasource.model;
+
+/** Enum for capturing the current datasource status. */
+public enum DataSourceStatus {
+  ACTIVE("active"),
+  DISABLED("disabled");
+
+  private String text;
+
+  DataSourceStatus(String text) {
+    this.text = text;
+  }
+
+  public String getText() {
+    return this.text;
+  }
+
+  /**
+   * Get DataSourceStatus from text.
+   *
+   * @param text text.
+   * @return DataSourceStatus {@link DataSourceStatus}.
+   */
+  public static DataSourceStatus fromString(String text) {
+    for (DataSourceStatus dataSourceStatus : DataSourceStatus.values()) {
+      if (dataSourceStatus.text.equalsIgnoreCase(text)) {
+        return dataSourceStatus;
+      }
+    }
+    throw new IllegalArgumentException("No DataSourceStatus with text " + text + " found");
+  }
+}