Skip to content

Commit 8604997

Browse files
ahmed-kamal2004ahmed-kamal2004
committed
extract node type using BFS
1 parent a2757ee commit 8604997

File tree

3 files changed

+153
-26
lines changed

3 files changed

+153
-26
lines changed

‎src/sqlancer/postgres/PostgresProvider.java

Lines changed: 26 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,15 @@
77
import java.sql.DriverManager;
88
import java.sql.SQLException;
99
import java.sql.Statement;
10+
import java.util.ArrayList;
1011
import java.util.Arrays;
12+
import java.util.LinkedList;
1113
import java.util.List;
14+
import java.util.Queue;
1215
import java.util.stream.Collectors;
1316

17+
import com.fasterxml.jackson.databind.JsonNode;
18+
import com.fasterxml.jackson.databind.ObjectMapper;
1419
import com.google.auto.service.AutoService;
1520

1621
import sqlancer.AbstractAction;
@@ -364,13 +369,6 @@ public String getQueryPlan(String selectStr, PostgresGlobalState globalState) th
364369
e.printStackTrace();
365370
}
366371
}
367-
368-
List<String> invalidList = List.of("create", "notify", "discard", "listen", "unlisten", "reset", "set",
369-
"delete", "alter", "analyze");
370-
if (invalidList.contains(selectStr.split("\\s+")[0].toLowerCase())) {
371-
return "";
372-
}
373-
374372
SQLQueryAdapter q = new SQLQueryAdapter(PostgresExplainGenerator.explain(selectStr), null);
375373
try (SQLancerResultSet rs = q.executeAndGet(globalState)) {
376374
while (rs.next()) {
@@ -404,27 +402,31 @@ protected boolean addRowsToAllTables(PostgresGlobalState globalState) throws Exc
404402
return true;
405403
}
406404

407-
private String formatQueryPlan(String queryPlan) {
408-
StringBuilder outQueryPlanFormatted = new StringBuilder();
409-
boolean insideBrackets = false;
405+
public String formatQueryPlan(String queryPlan) throws IOException {
406+
ObjectMapper mapper = new ObjectMapper();
407+
JsonNode root = mapper.readTree(queryPlan).get(0).get("Plan");
408+
// Extract nodes using BFS algorithm
409+
List<String> nodeTypes = extractNodeTypesIterative(root);
410+
return String.join(" ", nodeTypes);
411+
}
410412

411-
for (char ch : queryPlan.toCharArray()) {
412-
if (ch == '\n' || ch == ' ') {
413-
continue;
413+
// BFS algorithm for traversing the Json Query Plan
414+
private static List<String> extractNodeTypesIterative(JsonNode root) {
415+
List<String> result = new ArrayList<>();
416+
Queue<JsonNode> queue = new LinkedList<>();
417+
queue.add(root);
418+
while (!queue.isEmpty()) {
419+
JsonNode node = queue.poll();
420+
if (node.has("Node Type")) {
421+
result.add(node.get("Node Type").asText());
414422
}
415-
if (ch == '(') {
416-
insideBrackets = true;
417-
} else if (ch == ')') {
418-
insideBrackets = false;
419-
outQueryPlanFormatted.append(';');
420-
continue;
421-
}
422-
if (!insideBrackets) {
423-
outQueryPlanFormatted.append(ch);
423+
if (node.has("Plans") && node.get("Plans").isArray()) {
424+
for (JsonNode plan : node.get("Plans")) {
425+
queue.add(plan);
426+
}
424427
}
425428
}
426-
427-
return outQueryPlanFormatted.toString();
429+
return result;
428430
}
429431

430432
}

‎src/sqlancer/postgres/gen/PostgresExplainGenerator.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ private PostgresExplainGenerator() {
88

99
public static String explain(String selectStr) throws Exception {
1010
StringBuilder sb = new StringBuilder();
11-
sb.append("EXPLAIN ");
11+
sb.append("EXPLAIN (FORMAT JSON) ");
1212
sb.append(selectStr);
1313
return sb.toString();
1414
}

‎test/sqlancer/qpg/postgres/TestPostgresQueryPlan.java

Lines changed: 126 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,132 @@ void testPostgresQueryPlan() throws Exception {
3434
q = new SQLQueryAdapter("CREATE TABLE t2(c INT);", true);
3535
q.execute(state);
3636
String queryPlan = provider.getQueryPlan("SELECT * FROM t1 RIGHT JOIN t2 ON a<>0;", state);
37-
assertEquals("NestedLoopLeftJoin;->SeqScanont2;->Materialize;->SeqScanont1;Filter:;", queryPlan);
37+
assertEquals("Nested Loop Seq Scan Materialize Seq Scan", queryPlan);
38+
}
39+
40+
@Test
41+
void testFormatQueryPlan() throws Exception {
42+
43+
PostgresProvider provider = new PostgresProvider();
44+
45+
String queryPlan = "[\n" + " {\n" + " \"Plan\": {\n" + " \"Node Type\": \"Aggregate\",\n"
46+
+ " \"Strategy\": \"Hashed\",\n" + " \"Partial Mode\": \"Simple\",\n"
47+
+ " \"Parallel Aware\": false,\n" + " \"Async Capable\": false,\n"
48+
+ " \"Startup Cost\": 62998.82,\n" + " \"Total Cost\": 63009.32,\n"
49+
+ " \"Plan Rows\": 1050,\n" + " \"Plan Width\": 4,\n" + " \"Output\": [\"t1.c0\"],\n"
50+
+ " \"Group Key\": [\"t1.c0\"],\n" + " \"Planned Partitions\": 0,\n" + " \"Plans\": [\n"
51+
+ " {\n" + " \"Node Type\": \"Append\",\n"
52+
+ " \"Parent Relationship\": \"Outer\",\n" + " \"Parallel Aware\": false,\n"
53+
+ " \"Async Capable\": false,\n" + " \"Startup Cost\": 27150.40,\n"
54+
+ " \"Total Cost\": 62996.20,\n" + " \"Plan Rows\": 1050,\n"
55+
+ " \"Plan Width\": 4,\n" + " \"Subplans Removed\": 0,\n" + " \"Plans\": [\n"
56+
+ " {\n" + " \"Node Type\": \"Group\",\n"
57+
+ " \"Parent Relationship\": \"Member\",\n" + " \"Parallel Aware\": false,\n"
58+
+ " \"Async Capable\": false,\n" + " \"Startup Cost\": 27150.40,\n"
59+
+ " \"Total Cost\": 62949.08,\n" + " \"Plan Rows\": 200,\n"
60+
+ " \"Plan Width\": 4,\n" + " \"Output\": [\"t1.c0\"],\n"
61+
+ " \"Group Key\": [\"t1.c0\"],\n" + " \"Plans\": [\n" + " {\n"
62+
+ " \"Node Type\": \"Gather Merge\",\n"
63+
+ " \"Parent Relationship\": \"Outer\",\n"
64+
+ " \"Parallel Aware\": false,\n" + " \"Async Capable\": false,\n"
65+
+ " \"Startup Cost\": 27150.40,\n" + " \"Total Cost\": 62948.08,\n"
66+
+ " \"Plan Rows\": 400,\n" + " \"Plan Width\": 4,\n"
67+
+ " \"Output\": [\"t1.c0\"],\n" + " \"Workers Planned\": 2,\n"
68+
+ " \"Plans\": [\n" + " {\n"
69+
+ " \"Node Type\": \"Group\",\n"
70+
+ " \"Parent Relationship\": \"Outer\",\n"
71+
+ " \"Parallel Aware\": false,\n"
72+
+ " \"Async Capable\": false,\n"
73+
+ " \"Startup Cost\": 26150.38,\n"
74+
+ " \"Total Cost\": 61901.89,\n" + " \"Plan Rows\": 200,\n"
75+
+ " \"Plan Width\": 4,\n" + " \"Output\": [\"t1.c0\"],\n"
76+
+ " \"Group Key\": [\"t1.c0\"],\n" + " \"Plans\": [\n"
77+
+ " {\n" + " \"Node Type\": \"Merge Join\",\n"
78+
+ " \"Parent Relationship\": \"Outer\",\n"
79+
+ " \"Parallel Aware\": false,\n"
80+
+ " \"Async Capable\": false,\n"
81+
+ " \"Join Type\": \"Inner\",\n"
82+
+ " \"Startup Cost\": 26150.38,\n"
83+
+ " \"Total Cost\": 56906.48,\n"
84+
+ " \"Plan Rows\": 1998164,\n"
85+
+ " \"Plan Width\": 4,\n"
86+
+ " \"Output\": [\"t1.c0\"],\n"
87+
+ " \"Inner Unique\": false,\n"
88+
+ " \"Merge Cond\": \"(t0.c0 = t1.c0)\",\n"
89+
+ " \"Plans\": [\n" + " {\n"
90+
+ " \"Node Type\": \"Sort\",\n"
91+
+ " \"Parent Relationship\": \"Outer\",\n"
92+
+ " \"Parallel Aware\": false,\n"
93+
+ " \"Async Capable\": false,\n"
94+
+ " \"Startup Cost\": 25970.60,\n"
95+
+ " \"Total Cost\": 26362.39,\n"
96+
+ " \"Plan Rows\": 156719,\n"
97+
+ " \"Plan Width\": 4,\n"
98+
+ " \"Output\": [\"t0.c0\"],\n"
99+
+ " \"Sort Key\": [\"t0.c0\"],\n"
100+
+ " \"Plans\": [\n" + " {\n"
101+
+ " \"Node Type\": \"Seq Scan\",\n"
102+
+ " \"Parent Relationship\": \"Outer\",\n"
103+
+ " \"Parallel Aware\": true,\n"
104+
+ " \"Async Capable\": false,\n"
105+
+ " \"Relation Name\": \"t0\",\n"
106+
+ " \"Schema\": \"public\",\n"
107+
+ " \"Alias\": \"t0\",\n"
108+
+ " \"Startup Cost\": 0.00,\n"
109+
+ " \"Total Cost\": 10301.95,\n"
110+
+ " \"Plan Rows\": 156719,\n"
111+
+ " \"Plan Width\": 4,\n"
112+
+ " \"Output\": [\"t0.c0\"],\n"
113+
+ " \"Filter\": \"(t0.c0 < 100)\"\n"
114+
+ " }\n" + " ]\n"
115+
+ " },\n" + " {\n"
116+
+ " \"Node Type\": \"Sort\",\n"
117+
+ " \"Parent Relationship\": \"Inner\",\n"
118+
+ " \"Parallel Aware\": false,\n"
119+
+ " \"Async Capable\": false,\n"
120+
+ " \"Startup Cost\": 179.78,\n"
121+
+ " \"Total Cost\": 186.16,\n"
122+
+ " \"Plan Rows\": 2550,\n"
123+
+ " \"Plan Width\": 4,\n"
124+
+ " \"Output\": [\"t1.c0\"],\n"
125+
+ " \"Sort Key\": [\"t1.c0\"],\n"
126+
+ " \"Plans\": [\n" + " {\n"
127+
+ " \"Node Type\": \"Seq Scan\",\n"
128+
+ " \"Parent Relationship\": \"Outer\",\n"
129+
+ " \"Parallel Aware\": false,\n"
130+
+ " \"Async Capable\": false,\n"
131+
+ " \"Relation Name\": \"t1\",\n"
132+
+ " \"Schema\": \"public\",\n"
133+
+ " \"Alias\": \"t1\",\n"
134+
+ " \"Startup Cost\": 0.00,\n"
135+
+ " \"Total Cost\": 35.50,\n"
136+
+ " \"Plan Rows\": 2550,\n"
137+
+ " \"Plan Width\": 4,\n"
138+
+ " \"Output\": [\"t1.c0\"]\n" + " }\n"
139+
+ " ]\n" + " }\n"
140+
+ " ]\n" + " }\n" + " ]\n"
141+
+ " }\n" + " ]\n" + " }\n" + " ]\n"
142+
+ " },\n" + " {\n" + " \"Node Type\": \"Bitmap Heap Scan\",\n"
143+
+ " \"Parent Relationship\": \"Member\",\n" + " \"Parallel Aware\": false,\n"
144+
+ " \"Async Capable\": false,\n" + " \"Relation Name\": \"t2\",\n"
145+
+ " \"Schema\": \"public\",\n" + " \"Alias\": \"t2\",\n"
146+
+ " \"Startup Cost\": 10.74,\n" + " \"Total Cost\": 31.37,\n"
147+
+ " \"Plan Rows\": 850,\n" + " \"Plan Width\": 4,\n"
148+
+ " \"Output\": [\"t2.c0\"],\n" + " \"Recheck Cond\": \"(t2.c0 < 10)\",\n"
149+
+ " \"Plans\": [\n" + " {\n"
150+
+ " \"Node Type\": \"Bitmap Index Scan\",\n"
151+
+ " \"Parent Relationship\": \"Outer\",\n"
152+
+ " \"Parallel Aware\": false,\n" + " \"Async Capable\": false,\n"
153+
+ " \"Index Name\": \"t2_pkey\",\n" + " \"Startup Cost\": 0.00,\n"
154+
+ " \"Total Cost\": 10.53,\n" + " \"Plan Rows\": 850,\n"
155+
+ " \"Plan Width\": 0,\n" + " \"Index Cond\": \"(t2.c0 < 10)\"\n"
156+
+ " }\n" + " ]\n" + " }\n" + " ]\n" + " }\n"
157+
+ " ]\n" + " },\n" + " \"Planning Time\": 1.954\n" + " }\n" + "]\n";
158+
159+
String formatedQueryPlan = provider.formatQueryPlan(queryPlan);
160+
assertEquals(
161+
"Aggregate Append Group Bitmap Heap Scan Gather Merge Bitmap Index Scan Group Merge Join Sort Sort Seq Scan Seq Scan",
162+
formatedQueryPlan);
38163
}
39164

40165
}

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy