Skip to content

Commit b6d6fb1

Browse files
yingh0ngyingh0ng
committed
Merge remote-tracking branch 'origin/Fix-17026-import' into Fix-17026-import
2 parents a891208 + 69a8d6c commit b6d6fb1

File tree

2 files changed

+205
-1
lines changed

2 files changed

+205
-1
lines changed

‎dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-postgresql/src/main/java/org/apache/dolphinscheduler/plugin/datasource/postgresql/param/PostgreSQLDataSourceProcessor.java

+91-1
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,97 @@ public DataSourceProcessor create() {
132132
@Override
133133
public List<String> splitAndRemoveComment(String sql) {
134134
String cleanSQL = SQLParserUtils.removeComment(sql, com.alibaba.druid.DbType.postgresql);
135-
return SQLParserUtils.split(cleanSQL, com.alibaba.druid.DbType.postgresql);
135+
return splitSqlRespectingDollarQuotes(cleanSQL);
136+
}
137+
138+
private int findNextDollar(String sql, int from) {
139+
for (int i = from; i < sql.length(); i++) {
140+
if (sql.charAt(i) == '$') {
141+
return i;
142+
}
143+
if (!Character.isLetterOrDigit(sql.charAt(i))) {
144+
break; // Not a valid dollar tag
145+
}
146+
}
147+
return -1;
148+
}
149+
150+
private int findClosingTag(String sql, int startIndex, String tag) {
151+
boolean inString = false;
152+
153+
for (int i = startIndex; i <= sql.length() - tag.length(); i++) {
154+
char ch = sql.charAt(i);
155+
156+
if (ch == '\'') {
157+
// Handle escaped quote: ''
158+
if (i + 1 < sql.length() && sql.charAt(i + 1) == '\'') {
159+
i++; // skip escaped quote
160+
} else {
161+
inString = !inString;
162+
}
163+
}
164+
165+
if (!inString && sql.startsWith(tag, i)) {
166+
return i;
167+
}
168+
}
169+
170+
return -1;
171+
}
172+
173+
private List<String> splitSqlRespectingDollarQuotes(String sql) {
174+
List<String> result = new ArrayList<>();
175+
StringBuilder current = new StringBuilder();
176+
177+
boolean insideDollarBlock = false;
178+
String dollarTag = null;
179+
int i = 0;
180+
181+
while (i < sql.length()) {
182+
char ch = sql.charAt(i);
183+
184+
// Detect start of dollar-quote block (e.g. $DO$, $func$)
185+
if (!insideDollarBlock && ch == '$') {
186+
int tagEnd = findNextDollar(sql, i + 1);
187+
if (tagEnd > i) {
188+
String potentialTag = sql.substring(i, tagEnd + 1);
189+
int closingIndex = findClosingTag(sql, tagEnd + 1, potentialTag);
190+
if (closingIndex != -1) {
191+
// We're starting a dollar block
192+
insideDollarBlock = true;
193+
dollarTag = potentialTag;
194+
current.append(dollarTag);
195+
i = tagEnd + 1;
196+
continue;
197+
}
198+
}
199+
}
200+
201+
// Detect end of dollar-quote block
202+
if (insideDollarBlock && dollarTag != null && sql.startsWith(dollarTag, i)) {
203+
insideDollarBlock = false;
204+
current.append(dollarTag);
205+
i += dollarTag.length();
206+
continue;
207+
}
208+
209+
// Split only outside of a dollar block
210+
if (!insideDollarBlock && ch == ';') {
211+
result.add(current.toString().trim());
212+
current.setLength(0);
213+
i++;
214+
continue;
215+
}
216+
217+
current.append(ch);
218+
i++;
219+
}
220+
221+
if (!current.toString().trim().isEmpty()) {
222+
result.add(current.toString().trim());
223+
}
224+
225+
return result;
136226
}
137227

138228
private String transformOther(Map<String, String> otherMap) {

‎dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-postgresql/src/test/java/org/apache/dolphinscheduler/plugin/datasource/postgresql/PostgreSQLDataSourceProcessorTest.java

+114
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package org.apache.dolphinscheduler.plugin.datasource.postgresql;
19+
20+
import org.apache.dolphinscheduler.plugin.datasource.postgresql.param.PostgreSQLDataSourceProcessor;
21+
22+
import java.util.List;
23+
24+
import org.junit.jupiter.api.Assertions;
25+
import org.junit.jupiter.api.Test;
26+
27+
public class PostgreSQLDataSourceProcessorTest {
28+
29+
private final PostgreSQLDataSourceProcessor processor = new PostgreSQLDataSourceProcessor();
30+
31+
@Test
32+
public void testSingleDoBlock() {
33+
String sql = "DO $$ BEGIN RAISE NOTICE 'Hello'; END $$;";
34+
List<String> parts = processor.splitAndRemoveComment(sql);
35+
Assertions.assertEquals(1, parts.size());
36+
Assertions.assertEquals("DO $$ BEGIN RAISE NOTICE 'Hello'; END $$", parts.get(0));
37+
}
38+
39+
@Test
40+
public void testDoBlockAndInsert() {
41+
String sql = "DO $$ BEGIN RAISE NOTICE 'Start'; END $$;\nINSERT INTO test_table(id) VALUES (1);";
42+
List<String> parts = processor.splitAndRemoveComment(sql);
43+
Assertions.assertEquals(2, parts.size());
44+
}
45+
46+
@Test
47+
public void testStandardSql() {
48+
String sql = "INSERT INTO test_table(id) VALUES (1); UPDATE test_table SET id = 2 WHERE id = 1;";
49+
List<String> parts = processor.splitAndRemoveComment(sql);
50+
Assertions.assertEquals(2, parts.size());
51+
Assertions.assertEquals("INSERT INTO test_table(id) VALUES (1)", parts.get(0));
52+
Assertions.assertEquals("UPDATE test_table SET id = 2 WHERE id = 1", parts.get(1));
53+
}
54+
55+
@Test
56+
public void testCustomDollarTag() {
57+
String sql = "DO $func$ BEGIN RAISE NOTICE 'Tag test'; END $func$;";
58+
List<String> parts = processor.splitAndRemoveComment(sql);
59+
Assertions.assertEquals(1, parts.size());
60+
Assertions.assertTrue(parts.get(0).contains("$func$"));
61+
Assertions.assertEquals("DO $func$ BEGIN RAISE NOTICE 'Tag test'; END $func$", parts.get(0));
62+
}
63+
64+
@Test
65+
public void testCommentsPreserved() {
66+
String sql =
67+
"-- comment here\nDO $$ BEGIN NULL; END $$;\n-- trailing comment\nINSERT INTO test_table VALUES (5);";
68+
List<String> parts = processor.splitAndRemoveComment(sql);
69+
Assertions.assertEquals(2, parts.size());
70+
Assertions.assertEquals("DO $$ BEGIN NULL; END $$", parts.get(0));
71+
Assertions.assertEquals("INSERT INTO test_table VALUES (5)", parts.get(1));
72+
}
73+
74+
@Test
75+
public void testDoBlockWithSemicolonInsideString() {
76+
String sql = "DO $$ BEGIN RAISE NOTICE 'hello; world'; END $$;";
77+
List<String> parts = processor.splitAndRemoveComment(sql);
78+
Assertions.assertEquals(1, parts.size());
79+
Assertions.assertEquals("DO $$ BEGIN RAISE NOTICE 'hello; world'; END $$", parts.get(0));
80+
}
81+
82+
@Test
83+
public void testDoBlockWithDollarTagInsideString() {
84+
String sql = "DO $$ BEGIN RAISE NOTICE 'this has $DO$ inside'; END $$;";
85+
List<String> parts = processor.splitAndRemoveComment(sql);
86+
Assertions.assertEquals(1, parts.size());
87+
Assertions.assertEquals("DO $$ BEGIN RAISE NOTICE 'this has $DO$ inside'; END $$", parts.get(0));
88+
}
89+
90+
@Test
91+
public void testMultipleStatementsWithDoBlock() {
92+
String sql = "DO $$ BEGIN RAISE NOTICE 'msg'; END $$; INSERT INTO test_table VALUES (1);";
93+
List<String> parts = processor.splitAndRemoveComment(sql);
94+
Assertions.assertEquals(2, parts.size());
95+
Assertions.assertEquals("DO $$ BEGIN RAISE NOTICE 'msg'; END $$", parts.get(0));
96+
Assertions.assertEquals("INSERT INTO test_table VALUES (1)", parts.get(1));
97+
}
98+
99+
@Test
100+
public void testCustomDollarTagBlock() {
101+
String sql = "DO $func$ BEGIN RAISE NOTICE 'custom tag'; END $func$;";
102+
List<String> parts = processor.splitAndRemoveComment(sql);
103+
Assertions.assertEquals(1, parts.size());
104+
Assertions.assertEquals("DO $func$ BEGIN RAISE NOTICE 'custom tag'; END $func$", parts.get(0));
105+
}
106+
107+
@Test
108+
public void testDoBlockWithEscapedSingleQuote() {
109+
String sql = "DO $$ BEGIN RAISE NOTICE 'don''t split here'; END $$;";
110+
List<String> parts = processor.splitAndRemoveComment(sql);
111+
Assertions.assertEquals(1, parts.size());
112+
Assertions.assertEquals("DO $$ BEGIN RAISE NOTICE 'don''t split here'; END $$", parts.get(0));
113+
}
114+
}

0 commit comments

Comments
 (0)