Show column comment in JDBC based connector

trinodb · Nov 19, 2019 · 744c74c · 744c74c
1 parent d276843
commit 744c74c
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 0 deletions.
diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcClient.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcClient.java
@@ -56,6 +56,7 @@
 
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Strings.emptyToNull;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static com.google.common.base.Verify.verify;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
@@ -254,11 +255,14 @@ public List<JdbcColumnHandle> getColumns(ConnectorSession session, JdbcTableHand
                 // skip unsupported column types
                 if (columnMapping.isPresent()) {
                     boolean nullable = (resultSet.getInt("NULLABLE") != columnNoNulls);
+                    // Note: some databases (e.g. SQL Server) do not return column remarks/comment here.
+                    Optional<String> comment = Optional.ofNullable(emptyToNull(resultSet.getString("REMARKS")));
                     columns.add(JdbcColumnHandle.builder()
                             .setColumnName(columnName)
                             .setJdbcTypeHandle(typeHandle)
                             .setColumnType(columnMapping.get().getType())
                             .setNullable(nullable)
+                            .setComment(comment)
                             .build());
                 }
             }

diff --git a/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlIntegrationSmokeTest.java b/presto-mysql/src/test/java/io/prestosql/plugin/mysql/TestMySqlIntegrationSmokeTest.java
@@ -219,6 +219,19 @@ public void testInsertIntoNotNullColumn()
         assertUpdate("DROP TABLE test_insert_not_null");
     }
 
+    @Test
+    public void testColumnComment()
+            throws Exception
+    {
+        execute("CREATE TABLE tpch.test_column_comment (col1 bigint COMMENT 'test comment', col2 bigint COMMENT '', col3 bigint)");
+
+        assertQuery(
+                "SELECT column_name, comment FROM information_schema.columns WHERE table_schema = 'tpch' AND table_name = 'test_column_comment'",
+                "VALUES ('col1', 'test comment'), ('col2', null), ('col3', null)");
+
+        assertUpdate("DROP TABLE test_column_comment");
+    }
+
     private void execute(String sql)
             throws SQLException
     {

diff --git a/presto-postgresql/src/main/java/io/prestosql/plugin/postgresql/PostgreSqlClient.java b/presto-postgresql/src/main/java/io/prestosql/plugin/postgresql/PostgreSqlClient.java
@@ -238,11 +238,13 @@ public List<JdbcColumnHandle> getColumns(ConnectorSession session, JdbcTableHand
                     // skip unsupported column types
                     if (columnMapping.isPresent()) {
                         boolean nullable = (resultSet.getInt("NULLABLE") != columnNoNulls);
+                        Optional<String> comment = Optional.ofNullable(resultSet.getString("REMARKS"));
                         columns.add(JdbcColumnHandle.builder()
                                 .setColumnName(columnName)
                                 .setJdbcTypeHandle(typeHandle)
                                 .setColumnType(columnMapping.get().getType())
                                 .setNullable(nullable)
+                                .setComment(comment)
                                 .build());
                     }
                 }

diff --git a/...esql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlIntegrationSmokeTest.java b/...esql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlIntegrationSmokeTest.java
@@ -290,6 +290,21 @@ public void testInsertIntoNotNullColumn()
         assertUpdate("DROP TABLE test_insert_not_null");
     }
 
+    @Test
+    public void testColumnComment()
+            throws Exception
+    {
+        try (AutoCloseable ignoreTable = withTable("tpch.test_column_comment",
+                "(col1 bigint, col2 bigint, col3 bigint)")) {
+            execute("COMMENT ON COLUMN tpch.test_column_comment.col1 IS 'test comment'");
+            execute("COMMENT ON COLUMN tpch.test_column_comment.col2 IS ''"); // it will be NULL, PostgreSQL doesn't store empty comment
+
+            assertQuery(
+                    "SELECT column_name, comment FROM information_schema.columns WHERE table_schema = 'tpch' AND table_name = 'test_column_comment'",
+                    "VALUES ('col1', 'test comment'), ('col2', null), ('col3', null)");
+        }
+    }
+
     private AutoCloseable withSchema(String schema)
             throws Exception
     {

diff --git a/...server/src/test/java/io/prestosql/plugin/sqlserver/TestSqlServerIntegrationSmokeTest.java b/...server/src/test/java/io/prestosql/plugin/sqlserver/TestSqlServerIntegrationSmokeTest.java
@@ -74,6 +74,26 @@ public void testView()
         sqlServer.execute("DROP VIEW IF EXISTS test_view");
     }
 
+    @Test
+    public void testColumnComment()
+            throws Exception
+    {
+        try (AutoCloseable ignoreTable = withTable("test_column_comment",
+                "(col1 bigint, col2 bigint, col3 bigint)")) {
+            sqlServer.execute("" +
+                    "EXEC sp_addextendedproperty " +
+                    " 'MS_Description', 'test comment', " +
+                    " 'Schema', 'dbo', " +
+                    " 'Table', 'test_column_comment', " +
+                    " 'Column', 'col1'");
+
+            // SQL Server JDBC driver doesn't support REMARKS for column comment https://github.com/Microsoft/mssql-jdbc/issues/646
+            assertQuery(
+                    "SELECT column_name, comment FROM information_schema.columns WHERE table_schema = 'dbo' AND table_name = 'test_column_comment'",
+                    "VALUES ('col1', null), ('col2', null), ('col3', null)");
+        }
+    }
+
     @Test
     public void testDecimalPredicatePushdown()
             throws Exception