Skip to content

Commit

Permalink
Integrate into SparkQueryDispatcher
Browse files Browse the repository at this point in the history 
Signed-off-by: Tomoyuki Morita <moritato@amazon.com>
  • Loading branch information
@ykmr1224
ykmr1224 committed Sep 18, 2024
1 parent b3a543d commit 32ea4dd
Show file tree
Hide file tree
Showing 14 changed files with 988 additions and 817 deletions.
11 changes: 3 additions & 8 deletions async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
package org.opensearch.sql.spark.dispatcher;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import lombok.AllArgsConstructor;
import org.jetbrains.annotations.NotNull;
Expand All @@ -24,6 +23,7 @@
import org.opensearch.sql.spark.execution.session.SessionManager;
import org.opensearch.sql.spark.rest.model.LangType;
import org.opensearch.sql.spark.utils.SQLQueryUtils;
import org.opensearch.sql.spark.validator.SQLQueryValidator;

/** This class takes care of understanding query and dispatching job query to emr serverless. */
@AllArgsConstructor
Expand All @@ -38,6 +38,7 @@ public class SparkQueryDispatcher {
private final SessionManager sessionManager;
private final QueryHandlerFactory queryHandlerFactory;
private final QueryIdProvider queryIdProvider;
private final SQLQueryValidator sqlQueryValidator;

public DispatchQueryResponse dispatch(
DispatchQueryRequest dispatchQueryRequest,
Expand All @@ -54,13 +55,7 @@ public DispatchQueryResponse dispatch(
dispatchQueryRequest, asyncQueryRequestContext, dataSourceMetadata);
}

List<String> validationErrors =
SQLQueryUtils.validateSparkSqlQuery(
dataSourceService.getDataSource(dispatchQueryRequest.getDatasource()), query);
if (!validationErrors.isEmpty()) {
throw new IllegalArgumentException(
"Query is not allowed: " + String.join(", ", validationErrors));
}
sqlQueryValidator.validate(query, dataSourceMetadata.getConnector());
}
return handleDefaultQuery(dispatchQueryRequest, asyncQueryRequestContext, dataSourceMetadata);
}
Expand Down
8 changes: 8 additions & 0 deletions async-query-core/src/main/java/org/opensearch/sql/spark/utils/SQLQueryUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,14 @@ public static List<String> validateSparkSqlQuery(DataSource datasource, String s
}
}

public static SqlBaseParser getBaseParser(String sqlQuery) {
SqlBaseParser sqlBaseParser =
new SqlBaseParser(
new CommonTokenStream(new SqlBaseLexer(new CaseInsensitiveCharStream(sqlQuery))));
sqlBaseParser.addErrorListener(new SyntaxAnalysisErrorListener());
return sqlBaseParser;
}

private SqlBaseValidatorVisitor getSparkSqlValidatorVisitor(DataSource datasource) {
if (datasource != null
&& datasource.getConnectorType() != null
Expand Down
76 changes: 76 additions & 0 deletions ...c/main/java/org/opensearch/sql/spark/validator/CloudWatchLogsGrammarElementValidator.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.sql.spark.validator;

import static org.opensearch.sql.spark.validator.GrammarElement.*;

import com.google.common.collect.ImmutableSet;
import java.util.Set;

public class CloudWatchLogsGrammarElementValidator extends DenyListGrammarElementValidator {
private static final Set<GrammarElement> CWL_DENY_LIST =
ImmutableSet.<GrammarElement>builder()
.add(
ALTER_NAMESPACE,
ALTER_VIEW,
CREATE_NAMESPACE,
CREATE_FUNCTION,
CREATE_VIEW,
DROP_FUNCTION,
DROP_NAMESPACE,
DROP_VIEW,
REPAIR_TABLE,
TRUNCATE_TABLE,
INSERT,
LOAD,
EXPLAIN,
WITH,
CLUSTER_BY,
DISTRIBUTE_BY,
HINTS,
INLINE_TABLE,
FILE,
CROSS_JOIN,
LEFT_SEMI_JOIN,
RIGHT_OUTER_JOIN,
FULL_OUTER_JOIN,
LEFT_ANTI_JOIN,
TABLESAMPLE,
TABLE_VALUED_FUNCTION,
LATERAL_VIEW,
LATERAL_SUBQUERY,
TRANSFORM,
MANAGE_RESOURCE,
ANALYZE_TABLE,
CACHE_TABLE,
DESCRIBE_NAMESPACE,
DESCRIBE_FUNCTION,
DESCRIBE_QUERY,
DESCRIBE_TABLE,
REFRESH_RESOURCE,
REFRESH_TABLE,
REFRESH_FUNCTION,
RESET,
SET,
SHOW_COLUMNS,
SHOW_CREATE_TABLE,
SHOW_NAMESPACES,
SHOW_FUNCTIONS,
SHOW_PARTITIONS,
SHOW_TABLE_EXTENDED,
SHOW_TABLES,
SHOW_TBLPROPERTIES,
SHOW_VIEWS,
UNCACHE_TABLE,
CSV_FUNCTIONS,
MISC_FUNCTIONS,
UDF)
.build();

public CloudWatchLogsGrammarElementValidator() {
super(CWL_DENY_LIST);
}
}
10 changes: 5 additions & 5 deletions async-query-core/src/main/java/org/opensearch/sql/spark/validator/GrammarElement.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,16 @@

@AllArgsConstructor
enum GrammarElement {
ALTER_NAMESPACE("ALTER DATABASE/TABLE/NAMESPACE"),
ALTER_NAMESPACE("ALTER (DATABASE|TABLE|NAMESPACE)"),
ALTER_VIEW("ALTER VIEW"),
CREATE_NAMESPACE("CREATE DATABASE/TABLE/NAMESPACE"),
CREATE_NAMESPACE("CREATE (DATABASE|TABLE|NAMESPACE)"),
CREATE_FUNCTION("CREATE FUNCTION"),
CREATE_VIEW("CREATE VIEW"),
DROP_NAMESPACE("DROP DATABASE/TABLE/NAMESPACE"),
DROP_NAMESPACE("DROP (DATABASE|TABLE|NAMESPACE)"),
DROP_FUNCTION("DROP FUNCTION"),
DROP_VIEW("DROP VIEW"),
DROP_TABLE("DROP TABLE"),
REPAIR_TABLE("REPAIR TABLE"), // does this conflict with DROP_NAMESPACE?
REPAIR_TABLE("REPAIR TABLE"),
TRUNCATE_TABLE("TRUNCATE TABLE"),
// DML Statements
INSERT("INSERT"),
Expand Down Expand Up @@ -52,7 +52,7 @@ enum GrammarElement {
ANALYZE_TABLE("ANALYZE TABLE(S)"),
CACHE_TABLE("CACHE TABLE"),
CLEAR_CACHE("CLEAR CACHE"),
DESCRIBE_NAMESPACE("DESCRIBE (NAMESPACE|DATABASE|SCHEMA"),
DESCRIBE_NAMESPACE("DESCRIBE (NAMESPACE|DATABASE|SCHEMA)"),
DESCRIBE_FUNCTION("DESCRIBE FUNCTION"),
DESCRIBE_QUERY("DESCRIBE QUERY"),
DESCRIBE_TABLE("DESCRIBE TABLE"),
Expand Down
167 changes: 5 additions & 162 deletions ...core/src/main/java/org/opensearch/sql/spark/validator/GrammarElementValidatorFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,178 +5,21 @@

package org.opensearch.sql.spark.validator;

import static org.opensearch.sql.spark.validator.GrammarElement.*;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import java.util.Map;
import java.util.Set;
import org.opensearch.sql.datasource.model.DataSourceType;

public class GrammarElementValidatorFactory {
// Deny List for CloudWatch Logs datasource
private static final Set<GrammarElement> CWL_DENY_LIST =
ImmutableSet.<GrammarElement>builder()
.add(
ALTER_NAMESPACE,
ALTER_VIEW,
CREATE_NAMESPACE,
CREATE_FUNCTION,
CREATE_VIEW,
DROP_FUNCTION,
DROP_NAMESPACE,
DROP_VIEW,
REPAIR_TABLE,
TRUNCATE_TABLE,
INSERT,
LOAD,
EXPLAIN,
WITH,
CLUSTER_BY,
DISTRIBUTE_BY,
HINTS,
INLINE_TABLE,
FILE,
CROSS_JOIN,
LEFT_SEMI_JOIN,
RIGHT_OUTER_JOIN,
FULL_OUTER_JOIN,
LEFT_ANTI_JOIN,
TABLESAMPLE,
TABLE_VALUED_FUNCTION,
LATERAL_VIEW,
LATERAL_SUBQUERY,
TRANSFORM,
MANAGE_RESOURCE,
ANALYZE_TABLE,
CACHE_TABLE,
DESCRIBE_NAMESPACE,
DESCRIBE_FUNCTION,
DESCRIBE_QUERY,
DESCRIBE_TABLE,
REFRESH_RESOURCE,
REFRESH_TABLE,
REFRESH_FUNCTION,
RESET,
SET,
SHOW_COLUMNS,
SHOW_CREATE_TABLE,
SHOW_NAMESPACES,
SHOW_FUNCTIONS,
SHOW_PARTITIONS,
SHOW_TABLE_EXTENDED,
SHOW_TABLES,
SHOW_TBLPROPERTIES,
SHOW_VIEWS,
UNCACHE_TABLE,
CSV_FUNCTIONS,
MISC_FUNCTIONS,
UDF)
.build();

// Deny list for S3 Glue datasource
private static final Set<GrammarElement> S3GLUE_DENY_LIST =
ImmutableSet.<GrammarElement>builder()
.add(
ALTER_VIEW,
CREATE_FUNCTION,
CREATE_VIEW,
DROP_FUNCTION,
DROP_VIEW,
INSERT,
LOAD,
CLUSTER_BY,
DISTRIBUTE_BY,
HINTS,
INLINE_TABLE,
FILE,
CROSS_JOIN,
LEFT_SEMI_JOIN,
RIGHT_OUTER_JOIN,
FULL_OUTER_JOIN,
LEFT_ANTI_JOIN,
TABLESAMPLE,
TABLE_VALUED_FUNCTION,
TRANSFORM,
MANAGE_RESOURCE,
DESCRIBE_FUNCTION,
REFRESH_RESOURCE,
REFRESH_FUNCTION,
RESET,
SET,
SHOW_FUNCTIONS,
SHOW_VIEWS,
MISC_FUNCTIONS,
UDF)
.build();

// Deny list for Security Lake datasource
private static final Set<GrammarElement> SL_DENY_LIST =
ImmutableSet.<GrammarElement>builder()
.add(
ALTER_NAMESPACE,
ALTER_VIEW,
CREATE_NAMESPACE,
CREATE_FUNCTION,
CREATE_VIEW,
DROP_FUNCTION,
DROP_NAMESPACE,
DROP_VIEW,
REPAIR_TABLE,
TRUNCATE_TABLE,
INSERT,
LOAD,
CLUSTER_BY,
DISTRIBUTE_BY,
HINTS,
INLINE_TABLE,
FILE,
CROSS_JOIN,
LEFT_SEMI_JOIN,
RIGHT_OUTER_JOIN,
FULL_OUTER_JOIN,
LEFT_ANTI_JOIN,
TABLESAMPLE,
TABLE_VALUED_FUNCTION,
TRANSFORM,
MANAGE_RESOURCE,
ANALYZE_TABLE,
CACHE_TABLE,
CLEAR_CACHE,
DESCRIBE_NAMESPACE,
DESCRIBE_FUNCTION,
DESCRIBE_QUERY,
DESCRIBE_TABLE,
REFRESH_RESOURCE,
REFRESH_TABLE,
REFRESH_FUNCTION,
RESET,
SET,
SHOW_COLUMNS,
SHOW_CREATE_TABLE,
SHOW_NAMESPACES,
SHOW_FUNCTIONS,
SHOW_PARTITIONS,
SHOW_TABLE_EXTENDED,
SHOW_TABLES,
SHOW_TBLPROPERTIES,
SHOW_VIEWS,
UNCACHE_TABLE,
CSV_FUNCTIONS,
MISC_FUNCTIONS,
UDF)
.build();

private static GrammarElementValidator defaultValidator = new DenyListGrammarElementValidator(
ImmutableSet.of());
private static Map<DataSourceType, GrammarElementValidator> validatorMap =
ImmutableMap.of(
DataSourceType.S3GLUE, new DenyListGrammarElementValidator(S3GLUE_DENY_LIST),
DataSourceType.SECURITY_LAKE, new DenyListGrammarElementValidator(SL_DENY_LIST));
DataSourceType.S3GLUE, new S3GlueGrammarElementValidator(),
DataSourceType.SECURITY_LAKE, new SecurityLakeGrammarElementValidator());

public GrammarElementValidator getValidatorForDatasource(DataSourceType dataSourceType) {
return validatorMap.get(dataSourceType);
}

private static ImmutableSet.Builder<GrammarElement> copyBuilder(Set<GrammarElement> original) {
return ImmutableSet.<GrammarElement>builder().addAll(original);
return validatorMap.getOrDefault(dataSourceType, defaultValidator);
}
}
Loading

0 comments on commit 32ea4dd

Please sign in to comment.