Skip to content

Commit

Permalink
Add File support
Browse files Browse the repository at this point in the history 
Signed-off-by: Tomoyuki Morita <[email protected]>
  • Loading branch information
@ykmr1224
ykmr1224 committed Sep 18, 2024
1 parent 15ae225 commit b3a543d
Show file tree
Hide file tree
Showing 4 changed files with 47 additions and 26 deletions.
32 changes: 21 additions & 11 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 @@ -14,27 +14,29 @@
import org.opensearch.sql.datasource.model.DataSourceType;

public class GrammarElementValidatorFactory {
private static final Set<GrammarElement> DEFAULT_DENY_LIST =
ImmutableSet.of(CREATE_FUNCTION, DROP_FUNCTION, INSERT, LOAD, HINTS, TABLESAMPLE);

// Deny List for CloudWatch Logs datasource
private static final Set<GrammarElement> CWL_DENY_LIST =
copyBuilder(DEFAULT_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,
Expand Down Expand Up @@ -74,15 +76,20 @@ public class GrammarElementValidatorFactory {

// Deny list for S3 Glue datasource
private static final Set<GrammarElement> S3GLUE_DENY_LIST =
copyBuilder(DEFAULT_DENY_LIST)
ImmutableSet.<GrammarElement>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,
Expand All @@ -105,20 +112,25 @@ public class GrammarElementValidatorFactory {

// Deny list for Security Lake datasource
private static final Set<GrammarElement> SL_DENY_LIST =
copyBuilder(DEFAULT_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,
Expand Down Expand Up @@ -155,12 +167,10 @@ public class GrammarElementValidatorFactory {
UDF)
.build();


private static Map<DataSourceType, GrammarElementValidator> 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);
Expand Down
19 changes: 17 additions & 2 deletions async-query-core/src/main/java/org/opensearch/sql/spark/validator/SQLQueryValidator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -184,8 +185,6 @@ public Void visitAlterClusterBy(AlterClusterByContext ctx) {
return super.visitAlterClusterBy(ctx);
}



@Override
public Void visitSetNamespaceLocation(SetNamespaceLocationContext ctx) {
validateAllowed(GrammarElement.ALTER_NAMESPACE);
Expand Down Expand Up @@ -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);
Expand Down
1 change: 1 addition & 0 deletions async-query-core/src/test/java/org/opensearch/sql/spark/validator/FunctionTypeTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"));
}
}
21 changes: 8 additions & 13 deletions async-query-core/src/test/java/org/opensearch/sql/spark/validator/SQLQueryValidatorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -398,7 +402,6 @@ private static class VerifyValidator {

public void ok(TestQuery query) {
runValidate(validator, query.toString());

}

public void ng(TestQuery query) {
Expand All @@ -419,12 +422,4 @@ SingleStatementContext getParser(String query) {
return sqlBaseParser.singleStatement();
}
}

void ok(SQLQueryValidator validator, TestQuery query) {
}

void ng(SQLQueryValidator validator, TestQuery query) {
}


}

0 comments on commit b3a543d

Please sign in to comment.