Fix adding catalog property to session

trinodb · Sep 17, 2020 · e04cdde · e04cdde
1 parent b250c27
commit e04cdde
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 1 deletion.
diff --git a/presto-main/src/main/java/io/prestosql/Session.java b/presto-main/src/main/java/io/prestosql/Session.java
@@ -76,7 +76,9 @@ public final class Session
     private final ResourceEstimates resourceEstimates;
     private final Instant start;
     private final Map<String, String> systemProperties;
+    // TODO use Table
     private final Map<CatalogName, Map<String, String>> connectorProperties;
+    // TODO use Table
     private final Map<String, Map<String, String>> unprocessedCatalogProperties;
     private final SessionPropertyManager sessionPropertyManager;
     private final Map<String, String> preparedStatements;
@@ -574,7 +576,8 @@ private SessionBuilder(Session session)
             this.clientTags = ImmutableSet.copyOf(session.clientTags);
             this.start = session.start;
             this.systemProperties.putAll(session.systemProperties);
-            this.catalogSessionProperties.putAll(session.unprocessedCatalogProperties);
+            session.unprocessedCatalogProperties
+                    .forEach((catalog, properties) -> catalogSessionProperties.put(catalog, new HashMap<>(properties)));
             this.preparedStatements.putAll(session.preparedStatements);
         }
 

diff --git a/presto-main/src/test/java/io/prestosql/TestSession.java b/presto-main/src/test/java/io/prestosql/TestSession.java
@@ -62,4 +62,28 @@ public void testBuildWithCatalogProperty()
         assertThat(session.getConnectorProperties(new CatalogName("some_catalog")))
                 .isEqualTo(Map.of());
     }
+
+    @Test
+    public void testAddSecondCatalogProperty()
+    {
+        Session session = Session.builder(testSessionBuilder().build())
+                .setCatalogSessionProperty("some_catalog", "first_property", "some_value")
+                .build();
+        session = Session.builder(session)
+                .setCatalogSessionProperty("some_catalog", "second_property", "another_value")
+                .build();
+
+        assertThat(session.getUnprocessedCatalogProperties())
+                .isEqualTo(Map.of("some_catalog", Map.of(
+                        "first_property", "some_value",
+                        "second_property", "another_value")));
+
+        // empty, will be populated at transaction start
+        assertThat(session.getConnectorProperties())
+                .isEqualTo(Map.of());
+
+        // empty, will be populated at transaction start
+        assertThat(session.getConnectorProperties(new CatalogName("some_catalog")))
+                .isEqualTo(Map.of());
+    }
 }