Datasource disable feature

Signed-off-by: Vamsi Manohar <[email protected]>

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}.
@@ -84,4 +85,15 @@ public interface DataSourceService {
    * @param dataSourceName name of the {@link DataSource}.
    */
   Boolean dataSourceExists(String dataSourceName);
+
+  /**
+   * Performs authorization and datasource status check. We could have exposed the api using only
+   * datasource name, In order to avoid multiple calls form SparkQueryDispatcher, we are exposing
+   * API using {@link DataSourceMetadata}.
+   *
+   * <p>
+   *
+   * @param dataSourceMetadata name of the {@link DataSourceMetadata}
+   */
+  void verifyDataSourceAccess(DataSourceMetadata dataSourceMetadata);
 }

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;
@@ -57,96 +58,147 @@ 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 =
+        DataSourceStatus.ACTIVE; // Optional, default to false if not provided
+
+    public Builder setName(String name) {
+      this.name = name;
+      return this;
     }
-    if (resultIndex == null) {
-      this.resultIndex = fromNameToCustomResultIndex();
-    } else {
-      this.resultIndex = resultIndex;
+
+    public Builder setDescription(String description) {
+      this.description = description;
+      return this;
     }
 
-    this.connector = connector;
-    this.description = description;
-    this.properties = properties;
-    this.allowedRoles = allowedRoles;
-  }
+    public Builder setConnector(DataSourceType connector) {
+      this.connector = connector;
+      return this;
+    }
 
-  public DataSourceMetadata() {
-    this.description = StringUtils.EMPTY;
-    this.allowedRoles = new ArrayList<>();
-    this.properties = new HashMap<>();
-  }
+    public Builder setAllowedRoles(List<String> allowedRoles) {
+      this.allowedRoles = new ArrayList<>(allowedRoles);
+      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 setProperties(Map<String, String> properties) {
+      this.properties = new HashMap<>(properties);
+      return this;
+    }
 
-  public String validateCustomResultIndex(String resultIndex) {
-    if (resultIndex == null) {
-      return null;
+    public Builder setResultIndex(String resultIndex) {
+      this.resultIndex = resultIndex;
+      return this;
     }
-    if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) {
-      return INVALID_RESULT_INDEX_NAME_SIZE;
+
+    public Builder setDataSourceStatus(DataSourceStatus status) {
+      this.status = status;
+      return this;
     }
-    if (!resultIndex.matches(RESULT_INDEX_NAME_PATTERN)) {
-      return INVALID_CHAR_IN_RESULT_INDEX_NAME;
+
+    public DataSourceMetadata build() {
+      String errorMessage = validateCustomResultIndex(resultIndex);
+      if (errorMessage != null) {
+        throw new IllegalArgumentException(errorMessage);
+      }
+      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();
+      }
+      return new DataSourceMetadata(this);
     }
-    if (resultIndex != null && !resultIndex.startsWith(DEFAULT_RESULT_INDEX)) {
-      return INVALID_RESULT_INDEX_PREFIX;
+
+    public String validateCustomResultIndex(String resultIndex) {
+      if (resultIndex == null) {
+        return null;
+      }
+      if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) {
+        return INVALID_RESULT_INDEX_NAME_SIZE;
+      }
+      if (!resultIndex.matches(RESULT_INDEX_NAME_PATTERN)) {
+        return INVALID_CHAR_IN_RESULT_INDEX_NAME;
+      }
+      if (resultIndex != null && !resultIndex.startsWith(DEFAULT_RESULT_INDEX)) {
+        return INVALID_RESULT_INDEX_PREFIX;
+      }
+      return null;
     }
-    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);
+    /**
+     * 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();
     }
 
-    // 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);
+    public String fromNameToCustomResultIndex() {
+      if (name == null) {
+        throw new IllegalArgumentException("Datasource name cannot be null");
       }
+      return convertToValidResultIndex(DATASOURCE_TO_RESULT_INDEX.apply(name.toLowerCase()));
     }
-    return validChars.toString();
   }
 
-  public String fromNameToCustomResultIndex() {
-    if (name == null) {
-      throw new IllegalArgumentException("Datasource name cannot be null");
-    }
-    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,36 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.datasource.model;
+
+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");
+  }
+}

core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java

@@ -19,7 +19,6 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.opensearch.sql.DataSourceSchemaName;
 import org.opensearch.sql.analysis.symbol.Namespace;
@@ -196,13 +195,10 @@ public Set<DataSourceMetadata> getDataSourceMetadata(boolean isDefaultDataSource
       return Stream.of(opensearchDataSource, prometheusDataSource)
           .map(
               ds ->
-                  new DataSourceMetadata(
-                      ds.getName(),
-                      StringUtils.EMPTY,
-                      ds.getConnectorType(),
-                      Collections.emptyList(),
-                      ImmutableMap.of(),
-                      null))
+                  new DataSourceMetadata.Builder()
+                      .setName(ds.getName())
+                      .setConnector(ds.getConnectorType())
+                      .build())
           .collect(Collectors.toSet());
     }
 
@@ -243,6 +239,9 @@ public void deleteDataSource(String dataSourceName) {}
     public Boolean dataSourceExists(String dataSourceName) {
       return dataSourceName.equals(DEFAULT_DATASOURCE_NAME) || dataSourceName.equals("prometheus");
     }
+
+    @Override
+    public void verifyDataSourceAccess(DataSourceMetadata dataSourceMetadata) {}
   }
 
   private class TestTableFunctionImplementation implements TableFunctionImplementation {

core/src/test/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScanTest.java

@@ -13,12 +13,10 @@
 import static org.mockito.Mockito.when;
 
 import com.google.common.collect.ImmutableMap;
-import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.apache.commons.lang3.StringUtils;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -61,13 +59,10 @@ void testIterator() {
         dataSourceSet.stream()
             .map(
                 dataSource ->
-                    new DataSourceMetadata(
-                        dataSource.getName(),
-                        StringUtils.EMPTY,
-                        dataSource.getConnectorType(),
-                        Collections.emptyList(),
-                        ImmutableMap.of(),
-                        null))
+                    new DataSourceMetadata.Builder()
+                        .setName(dataSource.getName())
+                        .setConnector(dataSource.getConnectorType())
+                        .build())
             .collect(Collectors.toSet());
     when(dataSourceService.getDataSourceMetadata(false)).thenReturn(dataSourceMetadata);
 

datasources/src/main/java/org/opensearch/sql/datasources/exceptions/DatasourceDisabledException.java

@@ -0,0 +1,7 @@
+package org.opensearch.sql.datasources.exceptions;
+
+public class DatasourceDisabledException extends RuntimeException {
+  public DatasourceDisabledException(String message) {
+    super(message);
+  }
+}