Pass session to setup too

trinodb · Oct 1, 2020 · fea326b · fea326b
1 parent aa8dde6
commit fea326b
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 6 deletions.
diff --git a/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlTypeMapping.java b/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlTypeMapping.java
@@ -428,7 +428,7 @@ public void testDate()
                     .setTimeZoneKey(TimeZoneKey.getTimeZoneKey(timeZoneId))
                     .build();
             testCases.execute(getQueryRunner(), session, mysqlCreateAndInsert("tpch.test_date"));
-            testCases.execute(getQueryRunner(), session, prestoCreateAsSelect("test_date"));
+            testCases.execute(getQueryRunner(), session, prestoCreateAsSelect(session, "test_date"));
         }
     }
 
@@ -539,7 +539,12 @@ private void testUnsupportedDataType(String databaseDataType)
 
     private DataSetup prestoCreateAsSelect(String tableNamePrefix)
     {
-        return new CreateAsSelectDataSetup(new PrestoSqlExecutor(getQueryRunner()), tableNamePrefix);
+        return prestoCreateAsSelect(getSession(), tableNamePrefix);
+    }
+
+    private DataSetup prestoCreateAsSelect(Session session, String tableNamePrefix)
+    {
+        return new CreateAsSelectDataSetup(new PrestoSqlExecutor(getQueryRunner(), session), tableNamePrefix);
     }
 
     private DataSetup mysqlCreateAndInsert(String tableNamePrefix)

diff --git a/presto-phoenix/src/test/java/io/prestosql/plugin/phoenix/TestPhoenixSqlTypeMapping.java b/presto-phoenix/src/test/java/io/prestosql/plugin/phoenix/TestPhoenixSqlTypeMapping.java
@@ -236,7 +236,7 @@ public void testDate()
             Session session = Session.builder(getQueryRunner().getDefaultSession())
                     .setTimeZoneKey(TimeZoneKey.getTimeZoneKey(timeZoneId))
                     .build();
-            prestoTestCases.execute(getQueryRunner(), session, prestoCreateAsSelect("test_date"));
+            prestoTestCases.execute(getQueryRunner(), session, prestoCreateAsSelect(session, "test_date"));
             phoenixTestCases.execute(getQueryRunner(), session, phoenixCreateAndInsert("tpch.test_date"));
         }
     }
@@ -406,7 +406,12 @@ private DataType<Integer> primaryKey()
 
     private DataSetup prestoCreateAsSelect(String tableNamePrefix)
     {
-        return new CreateAsSelectDataSetup(new PrestoSqlExecutor(getQueryRunner()), tableNamePrefix);
+        return prestoCreateAsSelect(getSession(), tableNamePrefix);
+    }
+
+    private DataSetup prestoCreateAsSelect(Session session, String tableNamePrefix)
+    {
+        return new CreateAsSelectDataSetup(new PrestoSqlExecutor(getQueryRunner(), session), tableNamePrefix);
     }
 
     private DataSetup phoenixCreateAndInsert(String tableNamePrefix)

diff --git a/...to-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java b/...to-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java
@@ -944,7 +944,7 @@ public void testDate()
                     .setTimeZoneKey(TimeZoneKey.getTimeZoneKey(timeZoneId))
                     .build();
             testCases.execute(getQueryRunner(), session, postgresCreateAndInsert("tpch.test_date"));
-            testCases.execute(getQueryRunner(), session, prestoCreateAsSelect("test_date"));
+            testCases.execute(getQueryRunner(), session, prestoCreateAsSelect(session, "test_date"));
         }
     }
 
@@ -1039,7 +1039,7 @@ public void testTimestamp(boolean insertWithPresto, ZoneId sessionZone)
                 .build();
 
         if (insertWithPresto) {
-            tests.execute(getQueryRunner(), session, prestoCreateAsSelect("test_timestamp"));
+            tests.execute(getQueryRunner(), session, prestoCreateAsSelect(session, "test_timestamp"));
         }
         else {
             tests.execute(getQueryRunner(), session, postgresCreateAndInsert("tpch.test_timestamp"));