This is an automated email from the ASF dual-hosted git repository.
zihaoxiang pushed a commit to branch dev
in repository https://gitbox.apache.org/repos/asf/dolphinscheduler.git
The following commit(s) were added to refs/heads/dev by this push:
new 1caa65d6a1 [Fix][SQL][#17040] Preserve PostgreSQL PL/pgSQL DO $$
blocks during SQL splitting (#17145)
1caa65d6a1 is described below
commit 1caa65d6a168c1fc926800e23a1a3171e4823703
Author: Eshwaran <[email protected]>
AuthorDate: Wed Apr 23 05:35:24 2025 +0100
[Fix][SQL][#17040] Preserve PostgreSQL PL/pgSQL DO $$ blocks during SQL
splitting (#17145)
---
.../param/PostgreSQLDataSourceProcessor.java | 92 ++++++++++++++++-
.../PostgreSQLDataSourceProcessorTest.java | 114 +++++++++++++++++++++
2 files changed, 205 insertions(+), 1 deletion(-)
diff --git
a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-postgresql/src/main/java/org/apache/dolphinscheduler/plugin/datasource/postgresql/param/PostgreSQLDataSourceProcessor.java
b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-postgresql/src/main/java/org/apache/dolphinscheduler/plugin/datasource/postgresql/param/PostgreSQLDataSourceProcessor.java
index 260e5cb4af..b19ce008f0 100644
---
a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-postgresql/src/main/java/org/apache/dolphinscheduler/plugin/datasource/postgresql/param/PostgreSQLDataSourceProcessor.java
+++
b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-postgresql/src/main/java/org/apache/dolphinscheduler/plugin/datasource/postgresql/param/PostgreSQLDataSourceProcessor.java
@@ -132,7 +132,97 @@ public class PostgreSQLDataSourceProcessor extends
AbstractDataSourceProcessor {
@Override
public List<String> splitAndRemoveComment(String sql) {
String cleanSQL = SQLParserUtils.removeComment(sql,
com.alibaba.druid.DbType.postgresql);
- return SQLParserUtils.split(cleanSQL,
com.alibaba.druid.DbType.postgresql);
+ return splitSqlRespectingDollarQuotes(cleanSQL);
+ }
+
+ private int findNextDollar(String sql, int from) {
+ for (int i = from; i < sql.length(); i++) {
+ if (sql.charAt(i) == '$') {
+ return i;
+ }
+ if (!Character.isLetterOrDigit(sql.charAt(i))) {
+ break; // Not a valid dollar tag
+ }
+ }
+ return -1;
+ }
+
+ private int findClosingTag(String sql, int startIndex, String tag) {
+ boolean inString = false;
+
+ for (int i = startIndex; i <= sql.length() - tag.length(); i++) {
+ char ch = sql.charAt(i);
+
+ if (ch == '\'') {
+ // Handle escaped quote: ''
+ if (i + 1 < sql.length() && sql.charAt(i + 1) == '\'') {
+ i++; // skip escaped quote
+ } else {
+ inString = !inString;
+ }
+ }
+
+ if (!inString && sql.startsWith(tag, i)) {
+ return i;
+ }
+ }
+
+ return -1;
+ }
+
+ private List<String> splitSqlRespectingDollarQuotes(String sql) {
+ List<String> result = new ArrayList<>();
+ StringBuilder current = new StringBuilder();
+
+ boolean insideDollarBlock = false;
+ String dollarTag = null;
+ int i = 0;
+
+ while (i < sql.length()) {
+ char ch = sql.charAt(i);
+
+ // Detect start of dollar-quote block (e.g. $DO$, $func$)
+ if (!insideDollarBlock && ch == '$') {
+ int tagEnd = findNextDollar(sql, i + 1);
+ if (tagEnd > i) {
+ String potentialTag = sql.substring(i, tagEnd + 1);
+ int closingIndex = findClosingTag(sql, tagEnd + 1,
potentialTag);
+ if (closingIndex != -1) {
+ // We're starting a dollar block
+ insideDollarBlock = true;
+ dollarTag = potentialTag;
+ current.append(dollarTag);
+ i = tagEnd + 1;
+ continue;
+ }
+ }
+ }
+
+ // Detect end of dollar-quote block
+ if (insideDollarBlock && dollarTag != null &&
sql.startsWith(dollarTag, i)) {
+ insideDollarBlock = false;
+ current.append(dollarTag);
+ i += dollarTag.length();
+ continue;
+ }
+
+ // Split only outside of a dollar block
+ if (!insideDollarBlock && ch == ';') {
+ result.add(current.toString().trim());
+ current.setLength(0);
+ i++;
+ continue;
+ }
+
+ current.append(ch);
+ i++;
+ }
+
+ if (!current.toString().trim().isEmpty()) {
+ result.add(current.toString().trim());
+ }
+
+ return result;
}
private String transformOther(Map<String, String> otherMap) {
diff --git
a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-postgresql/src/test/java/org/apache/dolphinscheduler/plugin/datasource/postgresql/PostgreSQLDataSourceProcessorTest.java
b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-postgresql/src/test/java/org/apache/dolphinscheduler/plugin/datasource/postgresql/PostgreSQLDataSourceProcessorTest.java
new file mode 100644
index 0000000000..5be8d3d843
--- /dev/null
+++
b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-postgresql/src/test/java/org/apache/dolphinscheduler/plugin/datasource/postgresql/PostgreSQLDataSourceProcessorTest.java
@@ -0,0 +1,114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.dolphinscheduler.plugin.datasource.postgresql;
+
+import
org.apache.dolphinscheduler.plugin.datasource.postgresql.param.PostgreSQLDataSourceProcessor;
+
+import java.util.List;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class PostgreSQLDataSourceProcessorTest {
+
+ private final PostgreSQLDataSourceProcessor processor = new
PostgreSQLDataSourceProcessor();
+
+ @Test
+ public void testSingleDoBlock() {
+ String sql = "DO $$ BEGIN RAISE NOTICE 'Hello'; END $$;";
+ List<String> parts = processor.splitAndRemoveComment(sql);
+ Assertions.assertEquals(1, parts.size());
+ Assertions.assertEquals("DO $$ BEGIN RAISE NOTICE 'Hello'; END $$",
parts.get(0));
+ }
+
+ @Test
+ public void testDoBlockAndInsert() {
+ String sql = "DO $$ BEGIN RAISE NOTICE 'Start'; END $$;\nINSERT INTO
test_table(id) VALUES (1);";
+ List<String> parts = processor.splitAndRemoveComment(sql);
+ Assertions.assertEquals(2, parts.size());
+ }
+
+ @Test
+ public void testStandardSql() {
+ String sql = "INSERT INTO test_table(id) VALUES (1); UPDATE test_table
SET id = 2 WHERE id = 1;";
+ List<String> parts = processor.splitAndRemoveComment(sql);
+ Assertions.assertEquals(2, parts.size());
+ Assertions.assertEquals("INSERT INTO test_table(id) VALUES (1)",
parts.get(0));
+ Assertions.assertEquals("UPDATE test_table SET id = 2 WHERE id = 1",
parts.get(1));
+ }
+
+ @Test
+ public void testCustomDollarTag() {
+ String sql = "DO $func$ BEGIN RAISE NOTICE 'Tag test'; END $func$;";
+ List<String> parts = processor.splitAndRemoveComment(sql);
+ Assertions.assertEquals(1, parts.size());
+ Assertions.assertTrue(parts.get(0).contains("$func$"));
+ Assertions.assertEquals("DO $func$ BEGIN RAISE NOTICE 'Tag test'; END
$func$", parts.get(0));
+ }
+
+ @Test
+ public void testCommentsPreserved() {
+ String sql =
+ "-- comment here\nDO $$ BEGIN NULL; END $$;\n-- trailing
comment\nINSERT INTO test_table VALUES (5);";
+ List<String> parts = processor.splitAndRemoveComment(sql);
+ Assertions.assertEquals(2, parts.size());
+ Assertions.assertEquals("DO $$ BEGIN NULL; END $$", parts.get(0));
+ Assertions.assertEquals("INSERT INTO test_table VALUES (5)",
parts.get(1));
+ }
+
+ @Test
+ public void testDoBlockWithSemicolonInsideString() {
+ String sql = "DO $$ BEGIN RAISE NOTICE 'hello; world'; END $$;";
+ List<String> parts = processor.splitAndRemoveComment(sql);
+ Assertions.assertEquals(1, parts.size());
+ Assertions.assertEquals("DO $$ BEGIN RAISE NOTICE 'hello; world'; END
$$", parts.get(0));
+ }
+
+ @Test
+ public void testDoBlockWithDollarTagInsideString() {
+ String sql = "DO $$ BEGIN RAISE NOTICE 'this has $DO$ inside'; END
$$;";
+ List<String> parts = processor.splitAndRemoveComment(sql);
+ Assertions.assertEquals(1, parts.size());
+ Assertions.assertEquals("DO $$ BEGIN RAISE NOTICE 'this has $DO$
inside'; END $$", parts.get(0));
+ }
+
+ @Test
+ public void testMultipleStatementsWithDoBlock() {
+ String sql = "DO $$ BEGIN RAISE NOTICE 'msg'; END $$; INSERT INTO
test_table VALUES (1);";
+ List<String> parts = processor.splitAndRemoveComment(sql);
+ Assertions.assertEquals(2, parts.size());
+ Assertions.assertEquals("DO $$ BEGIN RAISE NOTICE 'msg'; END $$",
parts.get(0));
+ Assertions.assertEquals("INSERT INTO test_table VALUES (1)",
parts.get(1));
+ }
+
+ @Test
+ public void testCustomDollarTagBlock() {
+ String sql = "DO $func$ BEGIN RAISE NOTICE 'custom tag'; END $func$;";
+ List<String> parts = processor.splitAndRemoveComment(sql);
+ Assertions.assertEquals(1, parts.size());
+ Assertions.assertEquals("DO $func$ BEGIN RAISE NOTICE 'custom tag';
END $func$", parts.get(0));
+ }
+
+ @Test
+ public void testDoBlockWithEscapedSingleQuote() {
+ String sql = "DO $$ BEGIN RAISE NOTICE 'don''t split here'; END $$;";
+ List<String> parts = processor.splitAndRemoveComment(sql);
+ Assertions.assertEquals(1, parts.size());
+ Assertions.assertEquals("DO $$ BEGIN RAISE NOTICE 'don''t split here';
END $$", parts.get(0));
+ }
+}
