55
66package org .opensearch .sql .ppl .antlr ;
77
8- import static org .junit .Assert .assertEquals ;
9- import static org .junit .Assert .assertNotEquals ;
10- import static org .junit .Assert .assertNotNull ;
11- import static org .junit .Assert .assertThrows ;
12- import static org .junit .Assert .assertTrue ;
8+ import static org .junit .Assert .*;
139
1410import java .util .List ;
1511import org .antlr .v4 .runtime .CommonTokenStream ;
@@ -139,22 +135,22 @@ public void testDynamicSourceClauseParseTreeStructure() {
139135 searchFrom .fromClause ().dynamicSourceClause ();
140136
141137 // Verify source references size and text
142- assertEquals (
143- "Should have 2 source references" ,
144- 2 ,
145- dynamicSource .sourceReferences ().sourceReference ().size ());
138+ assertEquals ("Should have 2 source references" , 2 , dynamicSource .sourceReference ().size ());
146139 assertEquals (
147140 "Source references text should match" ,
148- "myindex,logs " ,
149- dynamicSource .sourceReferences ( ).getText ());
141+ "myindex" ,
142+ dynamicSource .sourceReference ( 0 ).getText ());
150143
151- // Verify filter args size and text
152144 assertEquals (
153- "Should have 2 filter args" , 2 , dynamicSource .sourceFilterArgs ().sourceFilterArg ().size ());
145+ "Source references text should match" , "logs" , dynamicSource .sourceReference (1 ).getText ());
146+ // Verify filter args size and text
147+ assertEquals ("Should have 2 filter args" , 2 , dynamicSource .sourceFilterArg ().size ());
154148 assertEquals (
155149 "Filter args text should match" ,
156- "fieldIndex=\" test\" ,count=100" ,
157- dynamicSource .sourceFilterArgs ().getText ());
150+ "fieldIndex=\" test\" " ,
151+ dynamicSource .sourceFilterArg (0 ).getText ());
152+ assertEquals (
153+ "Filter args text should match" , "count=100" , dynamicSource .sourceFilterArg (1 ).getText ());
158154 }
159155
160156 @ Test
@@ -172,22 +168,21 @@ public void testDynamicSourceWithComplexFilters() {
172168 searchFrom .fromClause ().dynamicSourceClause ();
173169
174170 // Verify source references
171+ assertEquals ("Should have 2 source references" , 2 , dynamicSource .sourceReference ().size ());
175172 assertEquals (
176- "Should have 2 source references" ,
177- 2 ,
178- dynamicSource .sourceReferences ().sourceReference ().size ());
173+ "First source reference text" , "vpc.flow_logs" , dynamicSource .sourceReference (0 ).getText ());
179174 assertEquals (
180- "Source references text" ,
181- "vpc.flow_logs,api.gateway" ,
182- dynamicSource .sourceReferences ().getText ());
175+ "Second source reference text" , "api.gateway" , dynamicSource .sourceReference (1 ).getText ());
183176
184177 // Verify filter args
178+ assertEquals ("Should have 3 filter args" , 3 , dynamicSource .sourceFilterArg ().size ());
185179 assertEquals (
186- "Should have 3 filter args" , 3 , dynamicSource .sourceFilterArgs ().sourceFilterArg ().size ());
180+ "First filter arg text" ,
181+ "region=\" us-east-1\" " ,
182+ dynamicSource .sourceFilterArg (0 ).getText ());
187183 assertEquals (
188- "Filter args text" ,
189- "region=\" us-east-1\" ,env=\" prod\" ,count=5" ,
190- dynamicSource .sourceFilterArgs ().getText ());
184+ "Second filter arg text" , "env=\" prod\" " , dynamicSource .sourceFilterArg (1 ).getText ());
185+ assertEquals ("Third filter arg text" , "count=5" , dynamicSource .sourceFilterArg (2 ).getText ());
191186 }
192187
193188 @ Test
@@ -203,24 +198,61 @@ public void testDynamicSourceWithSingleSource() {
203198 OpenSearchPPLParser .DynamicSourceClauseContext dynamicSource =
204199 searchFrom .fromClause ().dynamicSourceClause ();
205200
201+ assertEquals ("Should have 1 source reference" , 1 , dynamicSource .sourceReference ().size ());
202+ assertEquals ("Source reference text" , "ds:myindex" , dynamicSource .sourceReference (0 ).getText ());
203+
204+ assertEquals ("Should have 1 filter arg" , 1 , dynamicSource .sourceFilterArg ().size ());
206205 assertEquals (
207- "Should have 1 source reference" ,
208- 1 ,
209- dynamicSource .sourceReferences ().sourceReference ().size ());
210- assertEquals ("Source reference text" , "ds:myindex" , dynamicSource .sourceReferences ().getText ());
206+ "Filter arg text" , "fieldIndex=\" test\" " , dynamicSource .sourceFilterArg (0 ).getText ());
207+ }
208+
209+ @ Test
210+ public void testDynamicSourceWithoutSource () {
211+ String query = "source=[fieldIndex=\" httpStatus\" , region=\" us-west-2\" ]" ;
212+ OpenSearchPPLLexer lexer = new OpenSearchPPLLexer (new CaseInsensitiveCharStream (query ));
213+ OpenSearchPPLParser parser = new OpenSearchPPLParser (new CommonTokenStream (lexer ));
211214
215+ OpenSearchPPLParser .RootContext root = parser .root ();
216+ OpenSearchPPLParser .SearchFromContext searchFrom =
217+ (OpenSearchPPLParser .SearchFromContext )
218+ root .pplStatement ().queryStatement ().pplCommands ().searchCommand ();
219+ OpenSearchPPLParser .DynamicSourceClauseContext dynamicSource =
220+ searchFrom .fromClause ().dynamicSourceClause ();
221+ assertEquals ("Should have no source references" , 0 , dynamicSource .sourceReference ().size ());
222+ assertEquals ("Should have 2 filter arg" , 2 , dynamicSource .sourceFilterArg ().size ());
212223 assertEquals (
213- "Should have 1 filter arg" , 1 , dynamicSource .sourceFilterArgs ().sourceFilterArg ().size ());
224+ "First filter arg text" ,
225+ "fieldIndex=\" httpStatus\" " ,
226+ dynamicSource .sourceFilterArg (0 ).getText ());
214227 assertEquals (
215- "Filter arg text" , "fieldIndex=\" test\" " , dynamicSource .sourceFilterArgs ().getText ());
228+ "Second filter arg text" ,
229+ "region=\" us-west-2\" " ,
230+ dynamicSource .sourceFilterArg (1 ).getText ());
216231 }
217232
218233 @ Test
219- public void testDynamicSourceRequiresAtLeastOneSource () {
220- // The grammar requires at least one source reference before optional filter args
221- // This test documents that behavior
222- exceptionRule .expect (RuntimeException .class );
223- new PPLSyntaxParser ().parse ("source=[fieldIndex=\" httpStatus\" , region=\" us-west-2\" ]" );
234+ public void testDynamicSourceWithAllSources () {
235+ String query = "source=[`*`, fieldIndex=\" httpStatus\" , region=\" us-west-2\" ]" ;
236+ OpenSearchPPLLexer lexer = new OpenSearchPPLLexer (new CaseInsensitiveCharStream (query ));
237+ OpenSearchPPLParser parser = new OpenSearchPPLParser (new CommonTokenStream (lexer ));
238+
239+ OpenSearchPPLParser .RootContext root = parser .root ();
240+ OpenSearchPPLParser .SearchFromContext searchFrom =
241+ (OpenSearchPPLParser .SearchFromContext )
242+ root .pplStatement ().queryStatement ().pplCommands ().searchCommand ();
243+ OpenSearchPPLParser .DynamicSourceClauseContext dynamicSource =
244+ searchFrom .fromClause ().dynamicSourceClause ();
245+ assertEquals ("Should have 1 source reference" , 1 , dynamicSource .sourceReference ().size ());
246+ assertEquals ("Source reference text" , "`*`" , dynamicSource .sourceReference (0 ).getText ());
247+ assertEquals ("Should have 2 filter arg" , 2 , dynamicSource .sourceFilterArg ().size ());
248+ assertEquals (
249+ "First filter arg text" ,
250+ "fieldIndex=\" httpStatus\" " ,
251+ dynamicSource .sourceFilterArg (0 ).getText ());
252+ assertEquals (
253+ "Second filter arg text" ,
254+ "region=\" us-west-2\" " ,
255+ dynamicSource .sourceFilterArg (1 ).getText ());
224256 }
225257
226258 @ Test
@@ -236,18 +268,16 @@ public void testDynamicSourceWithDottedNames() {
236268 OpenSearchPPLParser .DynamicSourceClauseContext dynamicSource =
237269 searchFrom .fromClause ().dynamicSourceClause ();
238270
271+ assertEquals ("Should have 2 source references" , 2 , dynamicSource .sourceReference ().size ());
239272 assertEquals (
240- "Should have 2 source references" ,
241- 2 ,
242- dynamicSource .sourceReferences ().sourceReference ().size ());
273+ "First source reference text" , "vpc.flow_logs" , dynamicSource .sourceReference (0 ).getText ());
243274 assertEquals (
244- "Source references text" ,
245- "vpc.flow_logs, api.gateway.logs" ,
246- dynamicSource .sourceReferences ( ).getText ());
275+ "Second source reference text" ,
276+ "api.gateway.logs" ,
277+ dynamicSource .sourceReference ( 1 ).getText ());
247278
248- assertEquals (
249- "Should have 1 filter arg" , 1 , dynamicSource .sourceFilterArgs ().sourceFilterArg ().size ());
250- assertEquals ("Filter arg text" , "env=\" prod\" " , dynamicSource .sourceFilterArgs ().getText ());
279+ assertEquals ("Should have 1 filter arg" , 1 , dynamicSource .sourceFilterArg ().size ());
280+ assertEquals ("Filter arg text" , "env=\" prod\" " , dynamicSource .sourceFilterArg (0 ).getText ());
251281 }
252282
253283 @ Test
@@ -263,15 +293,11 @@ public void testDynamicSourceWithSimpleFilter() {
263293 OpenSearchPPLParser .DynamicSourceClauseContext dynamicSource =
264294 searchFrom .fromClause ().dynamicSourceClause ();
265295
266- assertEquals (
267- "Should have 1 source reference" ,
268- 1 ,
269- dynamicSource .sourceReferences ().sourceReference ().size ());
270- assertEquals ("Source reference text" , "logs" , dynamicSource .sourceReferences ().getText ());
296+ assertEquals ("Should have 1 source reference" , 1 , dynamicSource .sourceReference ().size ());
297+ assertEquals ("Source reference text" , "logs" , dynamicSource .sourceReference (0 ).getText ());
271298
272- assertEquals (
273- "Should have 1 filter arg" , 1 , dynamicSource .sourceFilterArgs ().sourceFilterArg ().size ());
274- assertEquals ("Filter arg text" , "count=100" , dynamicSource .sourceFilterArgs ().getText ());
299+ assertEquals ("Should have 1 filter arg" , 1 , dynamicSource .sourceFilterArg ().size ());
300+ assertEquals ("Filter arg text" , "count=100" , dynamicSource .sourceFilterArg (0 ).getText ());
275301 }
276302
277303 @ Test
@@ -292,12 +318,12 @@ public void testDynamicSourceWithInClause() {
292318 searchFrom .fromClause ().dynamicSourceClause ();
293319
294320 assertNotNull ("Dynamic source should exist" , dynamicSource );
295- assertNotNull ("Filter args should exist" , dynamicSource . sourceFilterArgs () );
321+ assertNotNull ("Filter args should exist" , dynamicSource );
296322
297323 // The IN clause should be parsed as a sourceFilterArg
298324 assertTrue (
299325 "Should have at least one filter arg with IN clause" ,
300- dynamicSource .sourceFilterArgs (). sourceFilterArg ().size () >= 1 );
326+ dynamicSource .sourceFilterArg ().size () >= 1 );
301327 }
302328
303329 @ Test
@@ -690,6 +716,124 @@ public void testBlockCommentShouldPass() {
690716 + " /* block comment */ " ));
691717 }
692718
719+ @ Test
720+ public void testDynamicSourceWithIntermixedSourcesAndFilters () {
721+ // Test intermixed sources and filters in various orders
722+ String query = "source=[myindex, region=\" us-east-1\" , logs, count=100, api.gateway]" ;
723+ OpenSearchPPLLexer lexer = new OpenSearchPPLLexer (new CaseInsensitiveCharStream (query ));
724+ OpenSearchPPLParser parser = new OpenSearchPPLParser (new CommonTokenStream (lexer ));
725+
726+ OpenSearchPPLParser .RootContext root = parser .root ();
727+ OpenSearchPPLParser .SearchFromContext searchFrom =
728+ (OpenSearchPPLParser .SearchFromContext )
729+ root .pplStatement ().queryStatement ().pplCommands ().searchCommand ();
730+ OpenSearchPPLParser .DynamicSourceClauseContext dynamicSource =
731+ searchFrom .fromClause ().dynamicSourceClause ();
732+
733+ // Verify we have 3 source references
734+ assertEquals ("Should have 3 source references" , 3 , dynamicSource .sourceReference ().size ());
735+ assertEquals ("First source reference" , "myindex" , dynamicSource .sourceReference (0 ).getText ());
736+ assertEquals ("Second source reference" , "logs" , dynamicSource .sourceReference (1 ).getText ());
737+ assertEquals (
738+ "Third source reference" , "api.gateway" , dynamicSource .sourceReference (2 ).getText ());
739+
740+ // Verify we have 2 filter args
741+ assertEquals ("Should have 2 filter args" , 2 , dynamicSource .sourceFilterArg ().size ());
742+ assertEquals (
743+ "First filter arg" , "region=\" us-east-1\" " , dynamicSource .sourceFilterArg (0 ).getText ());
744+ assertEquals ("Second filter arg" , "count=100" , dynamicSource .sourceFilterArg (1 ).getText ());
745+ }
746+
747+ @ Test
748+ public void testDynamicSourceStartingWithFilter () {
749+ // Test starting with a filter argument instead of source reference
750+ String query = "source=[region=\" us-west-1\" , myindex, env=\" prod\" , logs]" ;
751+ OpenSearchPPLLexer lexer = new OpenSearchPPLLexer (new CaseInsensitiveCharStream (query ));
752+ OpenSearchPPLParser parser = new OpenSearchPPLParser (new CommonTokenStream (lexer ));
753+
754+ OpenSearchPPLParser .RootContext root = parser .root ();
755+ OpenSearchPPLParser .SearchFromContext searchFrom =
756+ (OpenSearchPPLParser .SearchFromContext )
757+ root .pplStatement ().queryStatement ().pplCommands ().searchCommand ();
758+ OpenSearchPPLParser .DynamicSourceClauseContext dynamicSource =
759+ searchFrom .fromClause ().dynamicSourceClause ();
760+
761+ // Verify source references
762+ assertEquals ("Should have 2 source references" , 2 , dynamicSource .sourceReference ().size ());
763+ assertEquals ("First source reference" , "myindex" , dynamicSource .sourceReference (0 ).getText ());
764+ assertEquals ("Second source reference" , "logs" , dynamicSource .sourceReference (1 ).getText ());
765+
766+ // Verify filter args
767+ assertEquals ("Should have 2 filter args" , 2 , dynamicSource .sourceFilterArg ().size ());
768+ assertEquals (
769+ "First filter arg" , "region=\" us-west-1\" " , dynamicSource .sourceFilterArg (0 ).getText ());
770+ assertEquals ("Second filter arg" , "env=\" prod\" " , dynamicSource .sourceFilterArg (1 ).getText ());
771+ }
772+
773+ @ Test
774+ public void testDynamicSourceAlternatingPattern () {
775+ // Test alternating pattern of sources and filters
776+ String query =
777+ "source=[ds:index1, type=\" logs\" , ds:index2, count=50, ds:index3, region=\" us-east-1\" ]" ;
778+ OpenSearchPPLLexer lexer = new OpenSearchPPLLexer (new CaseInsensitiveCharStream (query ));
779+ OpenSearchPPLParser parser = new OpenSearchPPLParser (new CommonTokenStream (lexer ));
780+
781+ OpenSearchPPLParser .RootContext root = parser .root ();
782+ OpenSearchPPLParser .SearchFromContext searchFrom =
783+ (OpenSearchPPLParser .SearchFromContext )
784+ root .pplStatement ().queryStatement ().pplCommands ().searchCommand ();
785+ OpenSearchPPLParser .DynamicSourceClauseContext dynamicSource =
786+ searchFrom .fromClause ().dynamicSourceClause ();
787+
788+ // Verify source references
789+ assertEquals ("Should have 3 source references" , 3 , dynamicSource .sourceReference ().size ());
790+ assertEquals ("First source reference" , "ds:index1" , dynamicSource .sourceReference (0 ).getText ());
791+ assertEquals (
792+ "Second source reference" , "ds:index2" , dynamicSource .sourceReference (1 ).getText ());
793+ assertEquals ("Third source reference" , "ds:index3" , dynamicSource .sourceReference (2 ).getText ());
794+
795+ // Verify filter args
796+ assertEquals ("Should have 3 filter args" , 3 , dynamicSource .sourceFilterArg ().size ());
797+ assertEquals ("First filter arg" , "type=\" logs\" " , dynamicSource .sourceFilterArg (0 ).getText ());
798+ assertEquals ("Second filter arg" , "count=50" , dynamicSource .sourceFilterArg (1 ).getText ());
799+ assertEquals (
800+ "Third filter arg" , "region=\" us-east-1\" " , dynamicSource .sourceFilterArg (2 ).getText ());
801+ }
802+
803+ @ Test
804+ public void testDynamicSourceWithComplexIntermixedIN () {
805+ // Test intermixed with IN clause filter
806+ String query =
807+ "source=[logs, fieldIndex IN (\" status\" , \" error\" ), api.gateway, region=\" us-east-1\" ]" ;
808+ OpenSearchPPLLexer lexer = new OpenSearchPPLLexer (new CaseInsensitiveCharStream (query ));
809+ OpenSearchPPLParser parser = new OpenSearchPPLParser (new CommonTokenStream (lexer ));
810+
811+ OpenSearchPPLParser .RootContext root = parser .root ();
812+ assertNotNull ("Query should parse successfully" , root );
813+
814+ OpenSearchPPLParser .SearchFromContext searchFrom =
815+ (OpenSearchPPLParser .SearchFromContext )
816+ root .pplStatement ().queryStatement ().pplCommands ().searchCommand ();
817+ OpenSearchPPLParser .DynamicSourceClauseContext dynamicSource =
818+ searchFrom .fromClause ().dynamicSourceClause ();
819+
820+ assertNotNull ("Dynamic source should exist" , dynamicSource );
821+
822+ // Verify source references
823+ assertEquals ("Should have 2 source references" , 2 , dynamicSource .sourceReference ().size ());
824+ assertEquals ("First source reference" , "logs" , dynamicSource .sourceReference (0 ).getText ());
825+ assertEquals (
826+ "Second source reference" , "api.gateway" , dynamicSource .sourceReference (1 ).getText ());
827+
828+ // Verify filter args - should have IN clause and region filter
829+ assertEquals ("Should have 2 filter args" , 2 , dynamicSource .sourceFilterArg ().size ());
830+ assertTrue (
831+ "First filter should contain IN clause" ,
832+ dynamicSource .sourceFilterArg (0 ).getText ().contains ("IN" ));
833+ assertEquals (
834+ "Second filter arg" , "region=\" us-east-1\" " , dynamicSource .sourceFilterArg (1 ).getText ());
835+ }
836+
693837 @ Test
694838 public void testWhereCommand () {
695839 assertNotEquals (null , new PPLSyntaxParser ().parse ("SOURCE=test | WHERE x" ));
0 commit comments