feat: Add parameter binding to the Query tasks and Trigger triggers #376

(#493) add parameters binding keep legacy se of Statement for Query use PreparedStatement when using parameters
kestra-io · Jan 27, 2025 · aee7888 · aee7888
1 parent 7045b66
commit aee7888
Show file tree

Hide file tree

Showing 28 changed files with 186 additions and 76 deletions.
diff --git a/plugin-jdbc-arrow-flight/src/main/java/io/kestra/plugin/jdbc/arrowflight/Trigger.java b/plugin-jdbc-arrow-flight/src/main/java/io/kestra/plugin/jdbc/arrowflight/Trigger.java
@@ -73,6 +73,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-as400/src/main/java/io/kestra/plugin/jdbc/as400/Trigger.java b/plugin-jdbc-as400/src/main/java/io/kestra/plugin/jdbc/as400/Trigger.java
@@ -74,6 +74,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
 
         return query.run(runContext);

diff --git a/plugin-jdbc-clickhouse/src/main/java/io/kestra/plugin/jdbc/clickhouse/Trigger.java b/plugin-jdbc-clickhouse/src/main/java/io/kestra/plugin/jdbc/clickhouse/Trigger.java
@@ -72,6 +72,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-clickhouse/src/test/java/io/kestra/plugin/jdbc/clickhouse/ClickHouseTest.java b/plugin-jdbc-clickhouse/src/test/java/io/kestra/plugin/jdbc/clickhouse/ClickHouseTest.java
@@ -22,6 +22,7 @@
 import java.time.ZoneId;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 
 import static io.kestra.core.models.tasks.common.FetchType.FETCH;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -68,6 +69,26 @@ void select() throws Exception {
         assertThat(runOutput.getRow().get("Ipv6"), is("2a02:aa08:e000:3100:0:0:0:2"));
     }
 
+    @Test
+    void selectWithParameters() throws Exception {
+        RunContext runContext = runContextFactory.of(ImmutableMap.of());
+
+        Query task = Query.builder()
+            .url(Property.of(getUrl()))
+            .username(null)
+            .password(null)
+            .fetchType(Property.of(FetchType.FETCH_ONE))
+            .timeZoneId(Property.of("Europe/Paris"))
+            .parameters(Property.of(Map.of("num", 123)))
+            .sql(Property.of("select * from clickhouse_types where Int8 = :num"))
+            .build();
+
+        AbstractJdbcQuery.Output runOutput = task.run(runContext);
+        assertThat(runOutput.getRow(), notNullValue());
+
+        assertThat(runOutput.getRow().get("Int8"), is(123));
+    }
+
     @Test
     void update() throws Exception {
         RunContext runContext = runContextFactory.of(ImmutableMap.of());

diff --git a/plugin-jdbc-db2/src/main/java/io/kestra/plugin/jdbc/db2/Trigger.java b/plugin-jdbc-db2/src/main/java/io/kestra/plugin/jdbc/db2/Trigger.java
@@ -71,6 +71,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
 
         return query.run(runContext);

diff --git a/plugin-jdbc-dremio/src/main/java/io/kestra/plugin/jdbc/dremio/Trigger.java b/plugin-jdbc-dremio/src/main/java/io/kestra/plugin/jdbc/dremio/Trigger.java
@@ -72,6 +72,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-druid/src/main/java/io/kestra/plugin/jdbc/druid/Trigger.java b/plugin-jdbc-druid/src/main/java/io/kestra/plugin/jdbc/druid/Trigger.java
@@ -59,20 +59,21 @@ public class Trigger extends AbstractJdbcTrigger {
     @Override
     protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Exception {
         var query = Query.builder()
-                .id(this.id)
-                .type(Query.class.getName())
-                .url(this.getUrl())
-                .username(this.getUsername())
-                .password(this.getPassword())
-                .timeZoneId(this.getTimeZoneId())
-                .sql(this.getSql())
-                .fetch(this.isFetch())
-                .store(this.isStore())
-                .fetchOne(this.isFetchOne())
-                .fetchType(Property.of(this.renderFetchType(runContext)))
-                .fetchSize(this.getFetchSize())
-                .additionalVars(this.additionalVars)
-                .build();
+            .id(this.id)
+            .type(Query.class.getName())
+            .url(this.getUrl())
+            .username(this.getUsername())
+            .password(this.getPassword())
+            .timeZoneId(this.getTimeZoneId())
+            .sql(this.getSql())
+            .fetch(this.isFetch())
+            .store(this.isStore())
+            .fetchOne(this.isFetchOne())
+            .fetchType(Property.of(this.renderFetchType(runContext)))
+            .fetchSize(this.getFetchSize())
+            .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
+            .build();
         return query.run(runContext);
     }
 

diff --git a/plugin-jdbc-duckdb/src/main/java/io/kestra/plugin/jdbc/duckdb/Trigger.java b/plugin-jdbc-duckdb/src/main/java/io/kestra/plugin/jdbc/duckdb/Trigger.java
@@ -83,6 +83,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-duckdb/src/test/java/io/kestra/plugin/jdbc/duckdb/DuckDbTest.java b/plugin-jdbc-duckdb/src/test/java/io/kestra/plugin/jdbc/duckdb/DuckDbTest.java
@@ -36,7 +36,7 @@
  * - https://duckdb.org/docs/sql/data_types/overview
  */
 @KestraTest
-public class DuckDbTest {
+class DuckDbTest {
     @Inject
     private RunContextFactory runContextFactory;
 
@@ -153,7 +153,7 @@ void select() throws Exception {
         assertThat(runOutput.getRow().get("t_boolean"), is(true));
         assertThat(runOutput.getRow().get("t_date"), is(LocalDate.parse("1992-09-20")));
         assertThat(runOutput.getRow().get("t_double"), is(12345.12345D));
-        assertThat(runOutput.getRow().get("t_decimal"), is(BigDecimal.valueOf(12345123L, 3)));;
+        assertThat(runOutput.getRow().get("t_decimal"), is(BigDecimal.valueOf(12345123L, 3)));
         assertThat(runOutput.getRow().get("t_hugeint"), is("9223372036854775807"));
         assertThat(runOutput.getRow().get("t_integer"), is(2147483647));
         assertThat(runOutput.getRow().get("t_interval"), is("28 days"));
@@ -192,7 +192,7 @@ void selectFromExistingFileInUrl() throws Exception {
         assertThat(runOutput.getRow().get("t_boolean"), is(true));
         assertThat(runOutput.getRow().get("t_date"), is(LocalDate.parse("1992-09-20")));
         assertThat(runOutput.getRow().get("t_double"), is(12345.12345D));
-        assertThat(runOutput.getRow().get("t_decimal"), is(BigDecimal.valueOf(12345123L, 3)));;
+        assertThat(runOutput.getRow().get("t_decimal"), is(BigDecimal.valueOf(12345123L, 3)));
         assertThat(runOutput.getRow().get("t_hugeint"), is("9223372036854775807"));
         assertThat(runOutput.getRow().get("t_integer"), is(2147483647));
         assertThat(runOutput.getRow().get("t_interval"), is("28 days"));
@@ -209,6 +209,26 @@ void selectFromExistingFileInUrl() throws Exception {
         assertThat(runOutput.getRow().get("t_enum"), is("happy"));
     }
 
+    @Test
+    void selectFromExistingFileInUrlWithParameters() throws Exception {
+        RunContext runContext = runContextFactory.of(ImmutableMap.of());
+
+        URL resource = DuckDbTest.class.getClassLoader().getResource("db/duck.db");
+
+        Query task = Query.builder()
+            .fetchType(Property.of(FETCH_ONE))
+            .timeZoneId(Property.of("Europe/Paris"))
+            .url(Property.of("jdbc:duckdb:"+ Objects.requireNonNull(resource).getPath()))
+            .parameters(Property.of(Map.of("num", 2147483647)))
+            .sql(Property.of("SELECT * FROM duck_types WHERE t_integer = :num"))
+            .build();
+
+        AbstractJdbcQuery.Output runOutput = task.run(runContext);
+        assertThat(runOutput.getRow(), notNullValue());
+
+        assertThat(runOutput.getRow().get("t_integer"), is(2147483647));
+    }
+
 
     @Test
     void selectFromExistingFileInParameter() throws Exception {
@@ -232,7 +252,7 @@ void selectFromExistingFileInParameter() throws Exception {
         assertThat(runOutput.getRow().get("t_boolean"), is(true));
         assertThat(runOutput.getRow().get("t_date"), is(LocalDate.parse("1992-09-20")));
         assertThat(runOutput.getRow().get("t_double"), is(12345.12345D));
-        assertThat(runOutput.getRow().get("t_decimal"), is(BigDecimal.valueOf(12345123L, 3)));;
+        assertThat(runOutput.getRow().get("t_decimal"), is(BigDecimal.valueOf(12345123L, 3)));
         assertThat(runOutput.getRow().get("t_hugeint"), is("9223372036854775807"));
         assertThat(runOutput.getRow().get("t_integer"), is(2147483647));
         assertThat(runOutput.getRow().get("t_interval"), is("28 days"));

diff --git a/plugin-jdbc-mysql/src/main/java/io/kestra/plugin/jdbc/mysql/Trigger.java b/plugin-jdbc-mysql/src/main/java/io/kestra/plugin/jdbc/mysql/Trigger.java
@@ -77,6 +77,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-mysql/src/test/java/io/kestra/plugin/jdbc/mysql/MysqlTest.java b/plugin-jdbc-mysql/src/test/java/io/kestra/plugin/jdbc/mysql/MysqlTest.java
@@ -157,6 +157,26 @@ void update() throws Exception {
         assertThat(runOutput.getRow().get("d"), is("D"));
     }
 
+    @Test
+    void selectWithParameters() throws Exception {
+        RunContext runContext = runContextFactory.of(ImmutableMap.of());
+
+        Query task = Query.builder()
+            .url(Property.of(getUrl()))
+            .username(Property.of(getUsername()))
+            .password(Property.of(getPassword()))
+            .fetchType(Property.of(FETCH_ONE))
+            .timeZoneId(Property.of("Europe/Paris"))
+            .sql(Property.of("select * from mysql_types where concert_id=:concert_id"))
+            .parameters(Property.of(Map.of("concert_id", "1")))
+            .build();
+
+        AbstractJdbcQuery.Output runOutput = task.run(runContext);
+        assertThat(runOutput.getRow(), notNullValue());
+
+        assertThat(runOutput.getRow().get("concert_id"), is("1"));
+    }
+
     @Override
     protected String getUrl() {
         return "jdbc:mysql://127.0.0.1:64790/kestra";

diff --git a/plugin-jdbc-oracle/src/main/java/io/kestra/plugin/jdbc/oracle/Trigger.java b/plugin-jdbc-oracle/src/main/java/io/kestra/plugin/jdbc/oracle/Trigger.java
@@ -73,6 +73,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-oracle/src/test/java/io/kestra/plugin/jdbc/oracle/OracleTest.java b/plugin-jdbc-oracle/src/test/java/io/kestra/plugin/jdbc/oracle/OracleTest.java
@@ -18,6 +18,7 @@
 import java.time.LocalDate;
 import java.time.LocalDateTime;
 import java.time.ZonedDateTime;
+import java.util.Map;
 
 import static io.kestra.core.models.tasks.common.FetchType.FETCH_ONE;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -69,6 +70,26 @@ void select() throws Exception {
         assertThat(runOutput.getRow().get("T_TIMESTAMP_LOCAL"), is(LocalDateTime.parse("1998-01-23T12:00:00")));
     }
 
+    @Test
+    void selectWithParameters() throws Exception {
+        RunContext runContext = runContextFactory.of(ImmutableMap.of());
+
+        Query task = Query.builder()
+            .url(Property.of(getUrl()))
+            .username(Property.of(getUsername()))
+            .password(Property.of(getPassword()))
+            .fetchType(Property.of(FETCH_ONE))
+            .timeZoneId(Property.of("Europe/Paris"))
+            .parameters(Property.of(Map.of("char", "aa")))
+            .sql(Property.of("select * from oracle_types where T_CHAR = :char"))
+            .build();
+
+        AbstractJdbcQuery.Output runOutput = task.run(runContext);
+        assertThat(runOutput.getRow(), notNullValue());
+
+        assertThat(runOutput.getRow().get("T_CHAR"), is("aa"));
+    }
+
     @Test
     void update() throws Exception {
         RunContext runContext = runContextFactory.of(ImmutableMap.of());

diff --git a/plugin-jdbc-pinot/src/main/java/io/kestra/plugin/jdbc/pinot/Query.java b/plugin-jdbc-pinot/src/main/java/io/kestra/plugin/jdbc/pinot/Query.java
@@ -31,7 +31,7 @@
             code = """
                 id: pinot_query
                 namespace: company.team
-                
+
                 tasks:
                   - id: query
                     type: o.kestra.plugin.jdbc.pinot.Query

diff --git a/plugin-jdbc-postgres/src/main/java/io/kestra/plugin/jdbc/postgresql/Trigger.java b/plugin-jdbc-postgres/src/main/java/io/kestra/plugin/jdbc/postgresql/Trigger.java
@@ -77,6 +77,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-redshift/src/main/java/io/kestra/plugin/jdbc/redshift/Trigger.java b/plugin-jdbc-redshift/src/main/java/io/kestra/plugin/jdbc/redshift/Trigger.java
@@ -72,6 +72,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-snowflake/src/main/java/io/kestra/plugin/jdbc/snowflake/Trigger.java b/plugin-jdbc-snowflake/src/main/java/io/kestra/plugin/jdbc/snowflake/Trigger.java
@@ -83,6 +83,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .additionalVars(this.additionalVars)
             .warehouse(this.getWarehouse())
             .database(this.getDatabase())
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-sqlite/src/main/java/io/kestra/plugin/jdbc/sqlite/Trigger.java b/plugin-jdbc-sqlite/src/main/java/io/kestra/plugin/jdbc/sqlite/Trigger.java
@@ -68,6 +68,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .fetchSize(this.getFetchSize())
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-sqlserver/src/main/java/io/kestra/plugin/jdbc/sqlserver/Trigger.java b/plugin-jdbc-sqlserver/src/main/java/io/kestra/plugin/jdbc/sqlserver/Trigger.java
@@ -71,6 +71,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchOne(this.isFetchOne())
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-sybase/src/main/java/io/kestra/plugin/jdbc/sybase/Trigger.java b/plugin-jdbc-sybase/src/main/java/io/kestra/plugin/jdbc/sybase/Trigger.java
@@ -73,6 +73,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchOne(this.isFetchOne())
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
 
         return query.run(runContext);

diff --git a/plugin-jdbc-vectorwise/src/main/java/io/kestra/plugin/jdbc/vectorwise/Trigger.java b/plugin-jdbc-vectorwise/src/main/java/io/kestra/plugin/jdbc/vectorwise/Trigger.java
@@ -71,6 +71,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchOne(this.isFetchOne())
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }

diff --git a/plugin-jdbc-vertica/src/main/java/io/kestra/plugin/jdbc/vertica/Trigger.java b/plugin-jdbc-vertica/src/main/java/io/kestra/plugin/jdbc/vertica/Trigger.java
@@ -71,6 +71,7 @@ protected AbstractJdbcQuery.Output runQuery(RunContext runContext) throws Except
             .fetchOne(this.isFetchOne())
             .fetchType(Property.of(this.renderFetchType(runContext)))
             .additionalVars(this.additionalVars)
+            .parameters(this.getParameters())
             .build();
         return query.run(runContext);
     }