From 037a4cb9bde82f38e88b362dd920d8f3dc3bee22 Mon Sep 17 00:00:00 2001
From: praveenkrishna <praveenkrishna@tutanota.com>
Date: Fri, 25 Nov 2022 20:19:15 +0530
Subject: [PATCH] Propagate table information when creating Connection for
 RecordCursor

---
 .../src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java    | 2 +-
 .../src/main/java/io/trino/plugin/jdbc/CachingJdbcClient.java | 4 ++--
 .../main/java/io/trino/plugin/jdbc/ForwardingJdbcClient.java  | 4 ++--
 .../src/main/java/io/trino/plugin/jdbc/JdbcClient.java        | 2 +-
 .../src/main/java/io/trino/plugin/jdbc/JdbcRecordCursor.java  | 2 +-
 .../io/trino/plugin/jdbc/jmx/StatisticsAwareJdbcClient.java   | 4 ++--
 6 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java
index 151168859a8a..6bd84f8336ca 100644
--- a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java
+++ b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java
@@ -410,7 +410,7 @@ public ConnectorSplitSource getSplits(ConnectorSession session, JdbcTableHandle
     }
 
     @Override
-    public Connection getConnection(ConnectorSession session, JdbcSplit split)
+    public Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcTableHandle tableHandle)
             throws SQLException
     {
         Connection connection = connectionFactory.openConnection(session);
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/CachingJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/CachingJdbcClient.java
index 4a3d32daae3e..a905519ae4ab 100644
--- a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/CachingJdbcClient.java
+++ b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/CachingJdbcClient.java
@@ -208,10 +208,10 @@ public ConnectorSplitSource getSplits(ConnectorSession session, JdbcTableHandle
     }
 
     @Override
-    public Connection getConnection(ConnectorSession session, JdbcSplit split)
+    public Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcTableHandle tableHandle)
             throws SQLException
     {
-        return delegate.getConnection(session, split);
+        return delegate.getConnection(session, split, tableHandle);
     }
 
     @Override
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForwardingJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForwardingJdbcClient.java
index 84a772142479..230d08403d0b 100644
--- a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForwardingJdbcClient.java
+++ b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForwardingJdbcClient.java
@@ -139,10 +139,10 @@ public ConnectorSplitSource getSplits(ConnectorSession session, JdbcTableHandle
     }
 
     @Override
-    public Connection getConnection(ConnectorSession session, JdbcSplit split)
+    public Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcTableHandle tableHandle)
             throws SQLException
     {
-        return delegate().getConnection(session, split);
+        return delegate().getConnection(session, split, tableHandle);
     }
 
     @Override
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcClient.java
index fa5732a65971..0a33912a0f0c 100644
--- a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcClient.java
+++ b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcClient.java
@@ -85,7 +85,7 @@ default Optional<String> convertPredicate(ConnectorSession session, ConnectorExp
 
     ConnectorSplitSource getSplits(ConnectorSession session, JdbcTableHandle tableHandle);
 
-    Connection getConnection(ConnectorSession session, JdbcSplit split)
+    Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcTableHandle tableHandle)
             throws SQLException;
 
     default void abortReadConnection(Connection connection, ResultSet resultSet)
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcRecordCursor.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcRecordCursor.java
index 2768c4077594..c431b5efc2c8 100644
--- a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcRecordCursor.java
+++ b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcRecordCursor.java
@@ -78,7 +78,7 @@ public JdbcRecordCursor(JdbcClient jdbcClient, ExecutorService executor, Connect
         objectReadFunctions = new ObjectReadFunction[columnHandles.size()];
 
         try {
-            connection = jdbcClient.getConnection(session, split);
+            connection = jdbcClient.getConnection(session, split, table);
 
             for (int i = 0; i < this.columnHandles.length; i++) {
                 JdbcColumnHandle columnHandle = columnHandles.get(i);
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareJdbcClient.java
index 76d772b98109..aeff4b6d5360 100644
--- a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareJdbcClient.java
+++ b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareJdbcClient.java
@@ -159,10 +159,10 @@ public ConnectorSplitSource getSplits(ConnectorSession session, JdbcTableHandle
     }
 
     @Override
-    public Connection getConnection(ConnectorSession session, JdbcSplit split)
+    public Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcTableHandle tableHandle)
             throws SQLException
     {
-        return stats.getGetConnectionWithSplit().wrap(() -> delegate().getConnection(session, split));
+        return stats.getGetConnectionWithSplit().wrap(() -> delegate().getConnection(session, split, tableHandle));
     }
 
     @Override