Skip to content

Commit 2f2b0ca

Browse files
pratyakshsharmapratyakshsharma
committed
feat(connector): Add execute procedure in JDBC connectors
Cherry-pick of trinodb/trino#22556 Co-authored-by: ebyhr
1 parent b4a645a commit 2f2b0ca

File tree

23 files changed

+872
-7
lines changed

23 files changed

+872
-7
lines changed

‎presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcModule.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
*/
1414
package com.facebook.presto.plugin.jdbc;
1515

16+
import com.facebook.presto.plugin.jdbc.procedure.ExecuteProcedure;
1617
import com.facebook.presto.spi.connector.ConnectorAccessControl;
1718
import com.facebook.presto.spi.procedure.Procedure;
1819
import com.google.inject.Binder;
@@ -40,7 +41,7 @@ public JdbcModule(String connectorId)
4041
public void configure(Binder binder)
4142
{
4243
newOptionalBinder(binder, ConnectorAccessControl.class);
43-
newSetBinder(binder, Procedure.class);
44+
newSetBinder(binder, Procedure.class).addBinding().toProvider(ExecuteProcedure.class).in(Scopes.SINGLETON);
4445
binder.bind(JdbcConnectorId.class).toInstance(new JdbcConnectorId(connectorId));
4546

4647
binder.bind(JdbcMetadataCache.class).in(Scopes.SINGLETON);

‎presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/procedure/ExecuteProcedure.java‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.facebook.presto.plugin.jdbc.procedure;
15+
16+
import com.facebook.presto.plugin.jdbc.JdbcClient;
17+
import com.facebook.presto.plugin.jdbc.JdbcIdentity;
18+
import com.facebook.presto.plugin.jdbc.JdbcSplit;
19+
import com.facebook.presto.spi.ConnectorSession;
20+
import com.facebook.presto.spi.PrestoException;
21+
import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
22+
import com.facebook.presto.spi.procedure.Procedure;
23+
import com.facebook.presto.spi.procedure.Procedure.Argument;
24+
import com.google.common.collect.ImmutableList;
25+
import com.google.inject.Inject;
26+
import com.google.inject.Provider;
27+
28+
import java.lang.invoke.MethodHandle;
29+
import java.sql.Connection;
30+
import java.sql.SQLException;
31+
import java.sql.Statement;
32+
33+
import static com.facebook.presto.common.block.MethodHandleUtil.methodHandle;
34+
import static com.facebook.presto.common.type.StandardTypes.VARCHAR;
35+
import static com.facebook.presto.plugin.jdbc.JdbcErrorCode.JDBC_ERROR;
36+
import static com.google.common.base.MoreObjects.firstNonNull;
37+
import static java.util.Objects.requireNonNull;
38+
39+
public class ExecuteProcedure
40+
implements Provider<Procedure>
41+
{
42+
private static final MethodHandle EXECUTE = methodHandle(
43+
ExecuteProcedure.class,
44+
"execute",
45+
ConnectorSession.class,
46+
String.class);
47+
48+
private final JdbcClient jdbcClient;
49+
50+
@Inject
51+
public ExecuteProcedure(JdbcClient jdbcClient)
52+
{
53+
this.jdbcClient = requireNonNull(jdbcClient, "jdbcClient is null");
54+
}
55+
56+
public void execute(ConnectorSession session, String query)
57+
{
58+
try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(getClass().getClassLoader())) {
59+
doExecute(session, query);
60+
}
61+
}
62+
63+
private void doExecute(ConnectorSession session, String query)
64+
{
65+
try (Connection connection = jdbcClient.getConnection(session, JdbcIdentity.from(session), (JdbcSplit) null)) {
66+
connection.setReadOnly(false);
67+
try (Statement statement = connection.createStatement()) {
68+
//noinspection SqlSourceToSinkFlow
69+
statement.executeUpdate(query);
70+
}
71+
}
72+
catch (SQLException e) {
73+
throw new PrestoException(JDBC_ERROR, "Failed to execute query. " + firstNonNull(e.getMessage(), e), e);
74+
}
75+
}
76+
77+
@Override
78+
public Procedure get()
79+
{
80+
return new Procedure(
81+
"system",
82+
"execute",
83+
ImmutableList.of(new Argument("QUERY", VARCHAR)),
84+
EXECUTE.bindTo(this));
85+
}
86+
}

‎presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/JdbcQueryRunner.java‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import com.facebook.presto.tests.DistributedQueryRunner;
1818
import com.facebook.presto.tpch.TpchPlugin;
1919
import com.google.common.collect.ImmutableList;
20+
import com.google.common.collect.ImmutableMap;
2021
import io.airlift.tpch.TpchTable;
2122

2223
import java.sql.Connection;
@@ -41,10 +42,16 @@ private JdbcQueryRunner()
4142
public static DistributedQueryRunner createJdbcQueryRunner(TpchTable<?>... tables)
4243
throws Exception
4344
{
44-
return createJdbcQueryRunner(ImmutableList.copyOf(tables));
45+
return createJdbcQueryRunner(ImmutableMap.of(), ImmutableList.copyOf(tables));
4546
}
4647

47-
public static DistributedQueryRunner createJdbcQueryRunner(Iterable<TpchTable<?>> tables)
48+
public static DistributedQueryRunner createJdbcQueryRunner(Map<String, String> extraProperties, TpchTable<?>... tables)
49+
throws Exception
50+
{
51+
return createJdbcQueryRunner(extraProperties, ImmutableList.copyOf(tables));
52+
}
53+
54+
public static DistributedQueryRunner createJdbcQueryRunner(Map<String, String> extraProperties, Iterable<TpchTable<?>> tables)
4855
throws Exception
4956
{
5057
DistributedQueryRunner queryRunner = null;
@@ -55,10 +62,14 @@ public static DistributedQueryRunner createJdbcQueryRunner(Iterable<TpchTable<?>
5562
queryRunner.createCatalog("tpch", "tpch");
5663

5764
Map<String, String> properties = TestingH2JdbcModule.createProperties();
65+
Map<String, String> catalogProperties = ImmutableMap.<String, String>builder()
66+
.putAll(properties)
67+
.putAll(extraProperties)
68+
.build();
5869
createSchema(properties, "tpch");
5970

6071
queryRunner.installPlugin(new JdbcPlugin("base-jdbc", new TestingH2JdbcModule()));
61-
queryRunner.createCatalog("jdbc", "base-jdbc", properties);
72+
queryRunner.createCatalog("jdbc", "base-jdbc", catalogProperties);
6273

6374
copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, createSession(), tables);
6475

‎presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/TestJdbcClient.java‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,17 @@
1919
import com.facebook.presto.spi.ConnectorSession;
2020
import com.facebook.presto.spi.ConnectorTableMetadata;
2121
import com.facebook.presto.spi.SchemaTableName;
22+
import com.facebook.presto.testing.QueryRunner;
23+
import com.facebook.presto.tests.AbstractTestQueryFramework;
2224
import com.google.common.collect.ImmutableList;
2325
import com.google.common.collect.ImmutableSet;
2426
import org.testng.annotations.AfterClass;
2527
import org.testng.annotations.BeforeClass;
2628
import org.testng.annotations.Test;
2729

30+
import java.util.HashMap;
2831
import java.util.List;
32+
import java.util.Map;
2933
import java.util.Optional;
3034

3135
import static com.facebook.presto.common.type.BigintType.BIGINT;
@@ -34,6 +38,7 @@
3438
import static com.facebook.presto.common.type.RealType.REAL;
3539
import static com.facebook.presto.common.type.VarcharType.VARCHAR;
3640
import static com.facebook.presto.common.type.VarcharType.createVarcharType;
41+
import static com.facebook.presto.plugin.jdbc.JdbcQueryRunner.createJdbcQueryRunner;
3742
import static com.facebook.presto.plugin.jdbc.TestingDatabase.CONNECTOR_ID;
3843
import static com.facebook.presto.plugin.jdbc.TestingJdbcTypeHandle.JDBC_BIGINT;
3944
import static com.facebook.presto.plugin.jdbc.TestingJdbcTypeHandle.JDBC_DATE;
@@ -44,12 +49,14 @@
4449
import static java.util.Collections.emptyMap;
4550
import static java.util.Locale.ENGLISH;
4651
import static java.util.UUID.randomUUID;
52+
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
4753
import static org.testng.Assert.assertEquals;
4854
import static org.testng.Assert.assertNotNull;
4955
import static org.testng.Assert.assertTrue;
5056

5157
@Test
5258
public class TestJdbcClient
59+
extends AbstractTestQueryFramework
5360
{
5461
private static final ConnectorSession session = testSessionBuilder().build().toConnectorSession();
5562

@@ -189,4 +196,71 @@ public void testAlterColumns()
189196
jdbcClient.dropTable(session, JdbcIdentity.from(session), tableHandle);
190197
}
191198
}
199+
200+
@Test
201+
public void testExecuteProcedure()
202+
{
203+
String tableName = "test_execute";
204+
String schemaTableName = getSession().getSchema().orElseThrow() + "." + tableName;
205+
try {
206+
assertUpdate("CALL system.execute('CREATE TABLE " + schemaTableName + "(id BIGINT AUTO_INCREMENT PRIMARY KEY, a int)')");
207+
assertThat(getQueryRunner().tableExists(getSession(), tableName)).isTrue();
208+
assertUpdate("CALL system.execute('INSERT INTO " + schemaTableName + "(a) VALUES (1)')");
209+
assertUpdate("CALL system.execute('INSERT INTO " + schemaTableName + "(a) VALUES (21)')");
210+
assertQuery("SELECT * FROM " + schemaTableName, "VALUES (1, 1), (2, 21)");
211+
212+
assertUpdate("CALL system.execute('UPDATE " + schemaTableName + " SET a = 2 where id = 1')");
213+
assertQuery("SELECT * FROM " + schemaTableName, "VALUES (1, 2), (2, 21)");
214+
215+
assertUpdate("CALL system.execute('DELETE FROM " + schemaTableName + "')");
216+
assertQueryReturnsEmptyResult("SELECT * FROM " + schemaTableName);
217+
218+
assertUpdate("CALL system.execute('DROP TABLE " + schemaTableName + "')");
219+
assertThat(getQueryRunner().tableExists(getSession(), tableName)).isFalse();
220+
}
221+
finally {
222+
assertUpdate("DROP TABLE IF EXISTS " + schemaTableName);
223+
}
224+
}
225+
226+
@Test
227+
public void testExecuteProcedureWithNamedArgument()
228+
{
229+
String tableName = "test_execute_named";
230+
String schemaTableName = getSession().getSchema().orElseThrow() + "." + tableName;
231+
try {
232+
assertUpdate("CALL system.execute(QUERY => 'CREATE TABLE " + schemaTableName + "(id BIGINT AUTO_INCREMENT PRIMARY KEY, a int)')");
233+
assertThat(getQueryRunner().tableExists(getSession(), tableName)).isTrue();
234+
assertUpdate("CALL system.execute(QUERY => 'INSERT INTO " + schemaTableName + "(a) VALUES (1)')");
235+
assertUpdate("CALL system.execute(QUERY => 'INSERT INTO " + schemaTableName + "(a) VALUES (21)')");
236+
assertQuery("SELECT * FROM " + schemaTableName, "VALUES (1, 1), (2, 21)");
237+
238+
assertUpdate("CALL system.execute(QUERY => 'UPDATE " + schemaTableName + " SET a = 2 where id = 1')");
239+
assertQuery("SELECT * FROM " + schemaTableName, "VALUES (1, 2), (2, 21)");
240+
241+
assertUpdate("CALL system.execute(QUERY => 'DELETE FROM " + schemaTableName + "')");
242+
assertQueryReturnsEmptyResult("SELECT * FROM " + schemaTableName);
243+
244+
assertUpdate("CALL system.execute(QUERY => 'DROP TABLE " + schemaTableName + "')");
245+
assertThat(getQueryRunner().tableExists(getSession(), tableName)).isFalse();
246+
}
247+
finally {
248+
assertUpdate("DROP TABLE IF EXISTS " + schemaTableName);
249+
}
250+
}
251+
252+
@Test
253+
public void testExecuteProcedureWithInvalidQuery()
254+
{
255+
assertQueryFails("CALL system.execute('SELECT 1')", "(?s)Failed to execute query.*");
256+
assertQueryFails("CALL system.execute('invalid')", "(?s)Failed to execute query.*");
257+
}
258+
259+
@Override
260+
protected QueryRunner createQueryRunner() throws Exception
261+
{
262+
Map<String, String> properties = new HashMap<>();
263+
properties.put("allow-drop-table", "true");
264+
return createJdbcQueryRunner(properties);
265+
}
192266
}

‎presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/TestJdbcDistributedQueries.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import com.facebook.presto.Session;
1717
import com.facebook.presto.testing.QueryRunner;
1818
import com.facebook.presto.tests.AbstractTestQueries;
19+
import com.google.common.collect.ImmutableMap;
1920
import io.airlift.tpch.TpchTable;
2021
import org.testng.annotations.Test;
2122

@@ -29,7 +30,7 @@ public class TestJdbcDistributedQueries
2930
protected QueryRunner createQueryRunner()
3031
throws Exception
3132
{
32-
return createJdbcQueryRunner(TpchTable.getTables());
33+
return createJdbcQueryRunner(ImmutableMap.of(), TpchTable.getTables());
3334
}
3435

3536
@Override

‎presto-docs/src/main/sphinx/connector/hana.rst‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,38 @@ Property Name Description
9797
case-insensitively using lowercase normalization.
9898
================================================== ==================================================================== ===========
9999

100+
101+
Procedures
102+
----------
103+
104+
Use the :doc:`/sql/call` statement to perform data manipulation or administrative tasks. Procedures are available in the ``system`` schema of the catalog.
105+
106+
Execute Procedure
107+
^^^^^^^^^^^^^^^^^
108+
109+
Underlying datasource may support some operation or sql syntax which is not supported by Presto, either at the parser level or at the connector level.
110+
Trying to run such SQL statement via Presto can result in errors during parsing or analysing. For example, SAP Hana database supports creating auto generated
111+
primary keys which is not supported via presto. Running this procedure enables users to do a SQL passthrough to the underlying database and presto just acts
112+
as a middle man for passing the statement.
113+
114+
The following arguments are available:
115+
116+
============= ========== =============== =======================================================================
117+
Argument Name Required Type Description
118+
============= ========== =============== =======================================================================
119+
``QUERY`` Yes string Sql statement to run
120+
============= ========== =============== =======================================================================
121+
122+
Examples:
123+
124+
* Create a table with auto generated primary key::
125+
126+
CALL hana.system.execute('create table schema1.table1 (id BIGINT NOT NULL GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, a int)')
127+
128+
CALL hana.system.execute(QUERY => 'create table schema1.table1 (id BIGINT NOT NULL GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, a int)')
129+
130+
131+
100132
Querying HANA
101133
-------------------
102134

‎presto-docs/src/main/sphinx/connector/mysql.rst‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,36 @@ The connector maps PrestoDB types to the corresponding MySQL types:
221221

222222
No other types are supported.
223223

224+
Procedures
225+
----------
226+
227+
Use the :doc:`/sql/call` statement to perform data manipulation or administrative tasks. Procedures are available in the ``system`` schema of the catalog.
228+
229+
Execute Procedure
230+
^^^^^^^^^^^^^^^^^
231+
232+
Underlying datasource may support some operation or sql syntax which is not supported by Presto, either at the parser level or at the connector level.
233+
Trying to run such SQL statement via Presto can result in errors during parsing or analysing. For example, Mysql database supports creating AUTO_INCREMENT
234+
primary keys which is not supported via presto. Running this procedure enables users to do a SQL passthrough to the underlying database and presto just acts
235+
as a middle man for passing the statement.
236+
237+
The following arguments are available:
238+
239+
============= ========== =============== =======================================================================
240+
Argument Name Required Type Description
241+
============= ========== =============== =======================================================================
242+
``QUERY`` Yes string Sql statement to run
243+
============= ========== =============== =======================================================================
244+
245+
Examples:
246+
247+
* Create a table with AUTO_INCREMENT primary key::
248+
249+
CALL mysql.system.execute('create table schema1.table1 (id BIGINT AUTO_INCREMENT PRIMARY KEY, a int)')
250+
251+
CALL mysql.system.execute(QUERY => 'create table schema1.table1 (id BIGINT AUTO_INCREMENT PRIMARY KEY, a int)')
252+
253+
224254
SQL Support
225255
-----------
226256

‎presto-docs/src/main/sphinx/connector/oracle.rst‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,38 @@ Property Name Description
7676
case-insensitively using lowercase normalization.
7777
================================================== ==================================================================== ===========
7878

79+
80+
Procedures
81+
----------
82+
83+
Use the :doc:`/sql/call` statement to perform data manipulation or administrative tasks. Procedures are available in the ``system`` schema of the catalog.
84+
85+
Execute Procedure
86+
^^^^^^^^^^^^^^^^^
87+
88+
Underlying datasource may support some operation or sql syntax which is not supported by Presto, either at the parser level or at the connector level.
89+
Trying to run such SQL statement via Presto can result in errors during parsing or analysing. For example, Oracle database supports DELETE statement which is not
90+
supported via presto-oracle connector. Running this procedure enables users to do a SQL passthrough to the underlying database and presto just acts as a middle man
91+
for passing the statement.
92+
93+
The following arguments are available:
94+
95+
============= ========== =============== =======================================================================
96+
Argument Name Required Type Description
97+
============= ========== =============== =======================================================================
98+
``QUERY`` Yes string Sql statement to run
99+
============= ========== =============== =======================================================================
100+
101+
Examples:
102+
103+
* Delete a row from table `employees` where `employee_id = 101`::
104+
105+
CALL oracle.system.execute('delete from employees where employee_id = 101')
106+
107+
CALL oracle.system.execute(QUERY => 'delete from employees where employee_id = 101')
108+
109+
110+
79111
Querying Oracle
80112
---------------
81113

0 commit comments

Comments
 (0)