Skip to content

Commit c9611af

Browse files
shizy818shizy818
committed
add IT test for nest CTE
1 parent 49b7fa6 commit c9611af

File tree

3 files changed

+233
-4
lines changed

3 files changed

+233
-4
lines changed

‎integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,8 @@
5858
import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail;
5959
import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest;
6060
import static org.junit.Assert.assertEquals;
61+
import static org.junit.Assert.assertFalse;
62+
import static org.junit.Assert.assertTrue;
6163
import static org.junit.Assert.fail;
6264

6365
@RunWith(IoTDBTestRunner.class)
@@ -344,6 +346,60 @@ public void testNest() {
344346
testCteFailureWithVariants(cteTemplateQueries, mainQuery, errMsg);
345347
}
346348

349+
@Test
350+
public void testNestExplain1() throws SQLException {
351+
String sql =
352+
"explain with cte1 as (select * from testtb), "
353+
+ "cte2 as materialized (select time, voltage from cte1) "
354+
+ "select * from cte2";
355+
try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
356+
Statement statement = connection.createStatement()) {
357+
statement.execute("USE testdb");
358+
359+
// explain
360+
ResultSet resultSet = statement.executeQuery(sql);
361+
ResultSetMetaData metaData = resultSet.getMetaData();
362+
assertEquals(metaData.getColumnCount(), 1);
363+
assertEquals(metaData.getColumnName(1), "distribution plan");
364+
365+
StringBuilder sb = new StringBuilder();
366+
while (resultSet.next()) {
367+
sb.append(resultSet.getString(1)).append(System.lineSeparator());
368+
}
369+
String result = sb.toString();
370+
assertFalse(result.contains("CTE Query : 'cte1'"));
371+
assertTrue(result.contains("CTE Query : 'cte2'"));
372+
assertTrue(result.contains("Main Query"));
373+
}
374+
}
375+
376+
@Test
377+
public void testNestExplain2() throws SQLException {
378+
String sql =
379+
"explain with cte1 as materialized (select * from testtb), "
380+
+ "cte2 as materialized (select time, voltage from cte1) "
381+
+ "select * from cte2";
382+
try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
383+
Statement statement = connection.createStatement()) {
384+
statement.execute("USE testdb");
385+
386+
// explain
387+
ResultSet resultSet = statement.executeQuery(sql);
388+
ResultSetMetaData metaData = resultSet.getMetaData();
389+
assertEquals(metaData.getColumnCount(), 1);
390+
assertEquals(metaData.getColumnName(1), "distribution plan");
391+
392+
StringBuilder sb = new StringBuilder();
393+
while (resultSet.next()) {
394+
sb.append(resultSet.getString(1)).append(System.lineSeparator());
395+
}
396+
String result = sb.toString();
397+
assertTrue(result.contains("CTE Query : 'cte1'"));
398+
assertTrue(result.contains("CTE Query : 'cte2'"));
399+
assertTrue(result.contains("Main Query"));
400+
}
401+
}
402+
347403
@Test
348404
public void testRecursive() {
349405
String sqlTemplate =
@@ -395,7 +451,7 @@ public void testPrivileges() throws SQLException {
395451
}
396452
fail("No exception!");
397453
} catch (Exception e) {
398-
Assert.assertTrue(
454+
assertTrue(
399455
e.getMessage(),
400456
e.getMessage()
401457
.contains(

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteMaterializerTest.java‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -320,14 +320,19 @@ public void testExplainAnalyze() {
320320
assertPlan(logicalQueryPlan, output(explainAnalyze(cteScan)));
321321
}
322322

323-
/** This test primarily ensures code coverage for handleCteExplainResults method. */
323+
/**
324+
* This test primarily ensures code coverage: materializeCTE.handleCteExplainResults &
325+
* materializeCTE.fetchCteQueryResult
326+
*/
324327
@Test
325328
public void testExplain() {
326-
String sql = "with cte1 as materialized (SELECT time, s1 FROM table1) select * from cte1";
329+
String sql =
330+
"with cte1 as (select time, s1 from table1), "
331+
+ "cte2 as materialized (select * from cte1) select * from cte2";
327332

328333
LogicalQueryPlan logicalQueryPlan = planTester.createPlan(sql, true);
329334

330-
PlanMatchPattern cteScan = cteScan("cte1", ImmutableList.of("time", "s1"));
335+
PlanMatchPattern cteScan = cteScan("cte2", ImmutableList.of("time", "s1"));
331336

332337
// Verify full LogicalPlan
333338
/*

‎iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CteSubqueryTest.java‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
/*
2+
*
3+
* * Licensed to the Apache Software Foundation (ASF) under one
4+
* * or more contributor license agreements. See the NOTICE file
5+
* * distributed with this work for additional information
6+
* * regarding copyright ownership. The ASF licenses this file
7+
* * to you under the Apache License, Version 2.0 (the
8+
* * "License"); you may not use this file except in compliance
9+
* * with the License. You may obtain a copy of the License at
10+
* *
11+
* * http://www.apache.org/licenses/LICENSE-2.0
12+
* *
13+
* * Unless required by applicable law or agreed to in writing,
14+
* * software distributed under the License is distributed on an
15+
* * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
16+
* * KIND, either express or implied. See the License for the
17+
* * specific language governing permissions and limitations
18+
* * under the License.
19+
*
20+
*/
21+
22+
package org.apache.iotdb.db.queryengine.plan.relational.planner;
23+
24+
import org.apache.iotdb.common.rpc.thrift.TSStatus;
25+
import org.apache.iotdb.commons.exception.IoTDBException;
26+
import org.apache.iotdb.commons.schema.column.ColumnHeader;
27+
import org.apache.iotdb.db.protocol.session.SessionManager;
28+
import org.apache.iotdb.db.queryengine.common.QueryId;
29+
import org.apache.iotdb.db.queryengine.common.header.DatasetHeader;
30+
import org.apache.iotdb.db.queryengine.plan.Coordinator;
31+
import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
32+
import org.apache.iotdb.db.queryengine.plan.execution.QueryExecution;
33+
import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan;
34+
import org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern;
35+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ComparisonExpression;
36+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LongLiteral;
37+
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SymbolReference;
38+
39+
import com.google.common.collect.ImmutableList;
40+
import com.google.common.collect.ImmutableSet;
41+
import org.apache.tsfile.enums.TSDataType;
42+
import org.apache.tsfile.read.common.block.TsBlock;
43+
import org.apache.tsfile.read.common.block.column.LongColumn;
44+
import org.apache.tsfile.read.common.block.column.TimeColumn;
45+
import org.junit.Before;
46+
import org.junit.Test;
47+
import org.junit.runner.RunWith;
48+
import org.mockito.Mockito;
49+
import org.powermock.core.classloader.annotations.PowerMockIgnore;
50+
import org.powermock.core.classloader.annotations.PrepareForTest;
51+
import org.powermock.modules.junit4.PowerMockRunner;
52+
53+
import java.util.List;
54+
import java.util.Optional;
55+
56+
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanAssert.assertPlan;
57+
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.collect;
58+
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.exchange;
59+
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.output;
60+
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.tableScan;
61+
import static org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ComparisonExpression.Operator.EQUAL;
62+
import static org.powermock.api.mockito.PowerMockito.mockStatic;
63+
import static org.powermock.api.mockito.PowerMockito.when;
64+
65+
@PowerMockIgnore({"com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*", "javax.management.*"})
66+
@RunWith(PowerMockRunner.class)
67+
@PrepareForTest({Coordinator.class, SessionManager.class})
68+
public class CteSubqueryTest {
69+
private PlanTester planTester;
70+
71+
@Before
72+
public void setUp() throws Exception {
73+
planTester = new PlanTester();
74+
mockExecuteForTableModel();
75+
}
76+
77+
/**
78+
* This test primarily ensures code coverage:
79+
* PredicateWithUncorrelatedScalarSubqueryReconstructor.fetchUncorrelatedSubqueryResultForPredicate
80+
*/
81+
@Test
82+
public void testCteSubquery() throws IoTDBException {
83+
mockExecuteForTableModel();
84+
85+
String sql =
86+
"with cte1 as (select time, s1 from table1) select s1 from table1 "
87+
+ "where s1 = (select s2 from cte1)";
88+
89+
LogicalQueryPlan logicalQueryPlan = planTester.createPlan(sql);
90+
91+
PlanMatchPattern tableScan =
92+
tableScan(
93+
"testdb.table1",
94+
ImmutableList.of("s1"),
95+
ImmutableSet.of("s1"),
96+
new ComparisonExpression(EQUAL, new SymbolReference("s1"), new LongLiteral("1")));
97+
98+
// Verify full LogicalPlan
99+
/*
100+
* └──OutputNode
101+
* └──DeviceTableScanNode
102+
*/
103+
assertPlan(logicalQueryPlan, output(tableScan));
104+
105+
// Verify DistributionPlan
106+
assertPlan(planTester.getFragmentPlan(0), output(collect(exchange(), exchange(), exchange())));
107+
108+
assertPlan(planTester.getFragmentPlan(1), tableScan);
109+
assertPlan(planTester.getFragmentPlan(2), tableScan);
110+
assertPlan(planTester.getFragmentPlan(3), tableScan);
111+
}
112+
113+
private void mockExecuteForTableModel() throws IoTDBException {
114+
mockStatic(Coordinator.class);
115+
mockStatic(SessionManager.class);
116+
117+
// Create a mock Coordinator instance
118+
Coordinator mockCoordinator = Mockito.mock(Coordinator.class);
119+
when(Coordinator.getInstance()).thenReturn(mockCoordinator);
120+
121+
// Create mock SessionManager
122+
SessionManager mockSessionManager = Mockito.mock(SessionManager.class);
123+
when(SessionManager.getInstance()).thenReturn(mockSessionManager);
124+
125+
// Mock TSStatus with success status
126+
TSStatus mockStatus = Mockito.mock(TSStatus.class);
127+
when(mockStatus.getCode()).thenReturn(200); // Success status code
128+
129+
// Create a real ExecutionResult instance
130+
ExecutionResult mockResult = new ExecutionResult(new QueryId("1"), mockStatus);
131+
132+
// Mock the executeForTableModel method
133+
when(mockCoordinator.executeForTableModel(
134+
Mockito.any(), // Statement
135+
Mockito.any(), // SqlParser
136+
Mockito.any(), // IClientSession
137+
Mockito.anyLong(), // queryId
138+
Mockito.any(), // SessionInfo
139+
Mockito.anyString(), // String
140+
Mockito.any(), // Metadata
141+
Mockito.anyMap(), // Map<NodeRef<Table>, CteDataStore>
142+
Mockito.any(), // ExplainType
143+
Mockito.anyLong(), // timeOut
144+
Mockito.anyBoolean())) // userQuery
145+
.thenReturn(mockResult);
146+
147+
// Create QueryExecution mock
148+
QueryExecution mockQueryExecution = Mockito.mock(QueryExecution.class);
149+
when(mockQueryExecution.hasNextResult())
150+
.thenReturn(true) // First call returns true
151+
.thenReturn(false); // Subsequent calls return false
152+
153+
// Create a real DatasetHeader with time and s1 columns
154+
List<ColumnHeader> columnHeaders = ImmutableList.of(new ColumnHeader("s2", TSDataType.INT64));
155+
DatasetHeader mockDatasetHeader = new DatasetHeader(columnHeaders, false);
156+
when(mockQueryExecution.getDatasetHeader()).thenReturn(mockDatasetHeader);
157+
158+
// Create a TSBlock with sample data for getBatchResult
159+
160+
TimeColumn timeColumn = new TimeColumn(1, new long[] {1000L});
161+
LongColumn valueColumn = new LongColumn(1, Optional.empty(), new long[] {1L});
162+
TsBlock sampleTsBlock = new TsBlock(timeColumn, valueColumn);
163+
when(mockQueryExecution.getBatchResult()).thenReturn(Optional.of(sampleTsBlock));
164+
165+
// Mock coordinator methods
166+
when(mockCoordinator.getQueryExecution(Mockito.anyLong())).thenReturn(mockQueryExecution);
167+
}
168+
}

0 commit comments

Comments
 (0)