From b3a543de8204707b0288e1b35fa81c538ff58722 Mon Sep 17 00:00:00 2001 From: Tomoyuki Morita Date: Wed, 18 Sep 2024 12:28:00 -0700 Subject: [PATCH] Add File support Signed-off-by: Tomoyuki Morita --- .../GrammarElementValidatorFactory.java | 32 ++++++++++++------- .../spark/validator/SQLQueryValidator.java | 19 +++++++++-- .../sql/spark/validator/FunctionTypeTest.java | 1 + .../validator/SQLQueryValidatorTest.java | 21 +++++------- 4 files changed, 47 insertions(+), 26 deletions(-) diff --git a/async-query-core/src/main/java/org/opensearch/sql/spark/validator/GrammarElementValidatorFactory.java b/async-query-core/src/main/java/org/opensearch/sql/spark/validator/GrammarElementValidatorFactory.java index 71f989e456..ba2672e3df 100644 --- a/async-query-core/src/main/java/org/opensearch/sql/spark/validator/GrammarElementValidatorFactory.java +++ b/async-query-core/src/main/java/org/opensearch/sql/spark/validator/GrammarElementValidatorFactory.java @@ -14,27 +14,29 @@ import org.opensearch.sql.datasource.model.DataSourceType; public class GrammarElementValidatorFactory { - private static final Set DEFAULT_DENY_LIST = - ImmutableSet.of(CREATE_FUNCTION, DROP_FUNCTION, INSERT, LOAD, HINTS, TABLESAMPLE); - // Deny List for CloudWatch Logs datasource private static final Set CWL_DENY_LIST = - copyBuilder(DEFAULT_DENY_LIST) + ImmutableSet.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, @@ -74,15 +76,20 @@ public class GrammarElementValidatorFactory { // Deny list for S3 Glue datasource private static final Set S3GLUE_DENY_LIST = - copyBuilder(DEFAULT_DENY_LIST) + ImmutableSet.builder() .add( ALTER_VIEW, + CREATE_FUNCTION, CREATE_VIEW, + DROP_FUNCTION, DROP_VIEW, - DISTRIBUTE_BY, - INLINE_TABLE, + INSERT, + LOAD, CLUSTER_BY, DISTRIBUTE_BY, + HINTS, + INLINE_TABLE, + FILE, CROSS_JOIN, LEFT_SEMI_JOIN, RIGHT_OUTER_JOIN, @@ -105,20 +112,25 @@ public class GrammarElementValidatorFactory { // Deny list for Security Lake datasource private static final Set SL_DENY_LIST = - copyBuilder(DEFAULT_DENY_LIST) + ImmutableSet.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, @@ -155,12 +167,10 @@ public class GrammarElementValidatorFactory { UDF) .build(); - private static Map validatorMap = ImmutableMap.of( DataSourceType.S3GLUE, new DenyListGrammarElementValidator(S3GLUE_DENY_LIST), - DataSourceType.SECURITY_LAKE, new DenyListGrammarElementValidator(SL_DENY_LIST) - ); + DataSourceType.SECURITY_LAKE, new DenyListGrammarElementValidator(SL_DENY_LIST)); public GrammarElementValidator getValidatorForDatasource(DataSourceType dataSourceType) { return validatorMap.get(dataSourceType); diff --git a/async-query-core/src/main/java/org/opensearch/sql/spark/validator/SQLQueryValidator.java b/async-query-core/src/main/java/org/opensearch/sql/spark/validator/SQLQueryValidator.java index 14d7b1ce22..2fd91bf87a 100644 --- a/async-query-core/src/main/java/org/opensearch/sql/spark/validator/SQLQueryValidator.java +++ b/async-query-core/src/main/java/org/opensearch/sql/spark/validator/SQLQueryValidator.java @@ -79,6 +79,7 @@ import org.opensearch.sql.spark.antlr.parser.SqlBaseParser.ShowTablesContext; import org.opensearch.sql.spark.antlr.parser.SqlBaseParser.ShowTblPropertiesContext; import org.opensearch.sql.spark.antlr.parser.SqlBaseParser.ShowViewsContext; +import org.opensearch.sql.spark.antlr.parser.SqlBaseParser.TableNameContext; import org.opensearch.sql.spark.antlr.parser.SqlBaseParser.TableValuedFunctionContext; import org.opensearch.sql.spark.antlr.parser.SqlBaseParser.TransformClauseContext; import org.opensearch.sql.spark.antlr.parser.SqlBaseParser.TruncateTableContext; @@ -184,8 +185,6 @@ public Void visitAlterClusterBy(AlterClusterByContext ctx) { return super.visitAlterClusterBy(ctx); } - - @Override public Void visitSetNamespaceLocation(SetNamespaceLocationContext ctx) { validateAllowed(GrammarElement.ALTER_NAMESPACE); @@ -325,6 +324,22 @@ public Void visitExplain(ExplainContext ctx) { return super.visitExplain(ctx); } + @Override + public Void visitTableName(TableNameContext ctx) { + String reference = ctx.identifierReference().getText(); + System.out.println(reference); + if (isFileReference(reference)) { + validateAllowed(GrammarElement.FILE); + } + return super.visitTableName(ctx); + } + + private static final String FILE_REFERENCE_PATTERN = "^[a-zA-Z]+\\.`[^`]+`$"; + + private boolean isFileReference(String reference) { + return reference.matches(FILE_REFERENCE_PATTERN); + } + @Override public Void visitCtes(CtesContext ctx) { validateAllowed(GrammarElement.WITH); diff --git a/async-query-core/src/test/java/org/opensearch/sql/spark/validator/FunctionTypeTest.java b/async-query-core/src/test/java/org/opensearch/sql/spark/validator/FunctionTypeTest.java index 920d35df2f..a5f868421c 100644 --- a/async-query-core/src/test/java/org/opensearch/sql/spark/validator/FunctionTypeTest.java +++ b/async-query-core/src/test/java/org/opensearch/sql/spark/validator/FunctionTypeTest.java @@ -42,5 +42,6 @@ public void test() { assertEquals(FunctionType.MISC, FunctionType.fromFunctionName("version")); assertEquals(FunctionType.GENERATOR, FunctionType.fromFunctionName("explode")); assertEquals(FunctionType.GENERATOR, FunctionType.fromFunctionName("stack")); + assertEquals(FunctionType.UDF, FunctionType.fromFunctionName("unknown")); } } diff --git a/async-query-core/src/test/java/org/opensearch/sql/spark/validator/SQLQueryValidatorTest.java b/async-query-core/src/test/java/org/opensearch/sql/spark/validator/SQLQueryValidatorTest.java index 53a9a94a57..747485e718 100644 --- a/async-query-core/src/test/java/org/opensearch/sql/spark/validator/SQLQueryValidatorTest.java +++ b/async-query-core/src/test/java/org/opensearch/sql/spark/validator/SQLQueryValidatorTest.java @@ -164,7 +164,9 @@ public String toString() { @Test void s3glueQueries() { - VerifyValidator v = new VerifyValidator(new SQLQueryValidator(factory.getValidatorForDatasource(DataSourceType.S3GLUE))); + VerifyValidator v = + new VerifyValidator( + new SQLQueryValidator(factory.getValidatorForDatasource(DataSourceType.S3GLUE))); // DDL Statements v.ok(TestQuery.ALTER_DATABASE); v.ok(TestQuery.ALTER_TABLE); @@ -195,7 +197,7 @@ void s3glueQueries() { v.ok(TestQuery.HAVING_CLAUSE); v.ng(TestQuery.HINTS); v.ng(TestQuery.INLINE_TABLE); - // v.ng(TestQuery.FILE); TODO: need dive deep + v.ng(TestQuery.FILE); v.ok(TestQuery.INNER_JOIN); v.ng(TestQuery.CROSS_JOIN); v.ok(TestQuery.LEFT_OUTER_JOIN); @@ -279,7 +281,9 @@ void s3glueQueries() { @Test void securityLakeQueries() { - VerifyValidator v = new VerifyValidator(new SQLQueryValidator(factory.getValidatorForDatasource(DataSourceType.SECURITY_LAKE))); + VerifyValidator v = + new VerifyValidator( + new SQLQueryValidator(factory.getValidatorForDatasource(DataSourceType.SECURITY_LAKE))); // DDL Statements v.ng(TestQuery.ALTER_DATABASE); v.ng(TestQuery.ALTER_TABLE); @@ -310,7 +314,7 @@ void securityLakeQueries() { v.ok(TestQuery.HAVING_CLAUSE); v.ng(TestQuery.HINTS); v.ng(TestQuery.INLINE_TABLE); - // v.ng(TestQuery.FILE); TODO: need dive deep + v.ng(TestQuery.FILE); v.ok(TestQuery.INNER_JOIN); v.ng(TestQuery.CROSS_JOIN); v.ok(TestQuery.LEFT_OUTER_JOIN); @@ -398,7 +402,6 @@ private static class VerifyValidator { public void ok(TestQuery query) { runValidate(validator, query.toString()); - } public void ng(TestQuery query) { @@ -419,12 +422,4 @@ SingleStatementContext getParser(String query) { return sqlBaseParser.singleStatement(); } } - - void ok(SQLQueryValidator validator, TestQuery query) { - } - - void ng(SQLQueryValidator validator, TestQuery query) { - } - - }