Return column comments in SQLAlchemy get_columns

tests/integration/test_sqlalchemy_integration.py

@@ -671,6 +671,34 @@ def test_get_table_comment(trino_connection):
         metadata.drop_all(engine)
 
 
+@pytest.mark.skipif(
+    trino_version() < 395,
+    reason="Memory connector supports column comments since Trino 395"
+)
+@pytest.mark.parametrize('trino_connection', ['memory'], indirect=True)
+def test_get_columns_with_comments(trino_connection):
+    engine, conn = trino_connection
+
+    if not engine.dialect.has_schema(conn, "test"):
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema("test"))
+
+    table_name = "test.table_with_column_comments"
+    try:
+        conn.execute(sqla.text(
+            f"CREATE TABLE {table_name} ("
+            "id INTEGER COMMENT 'this is the id column', "
+            "name VARCHAR"
+            ")"
+        ))
+        columns = sqla.inspect(engine).get_columns(table_name='table_with_column_comments', schema="test")
+        columns_by_name = {column['name']: column for column in columns}
+        assert columns_by_name['id']['comment'] == 'this is the id column'
+        assert columns_by_name['name']['comment'] is None
+    finally:
+        conn.execute(sqla.text(f"DROP TABLE IF EXISTS {table_name}"))
+
+
 @pytest.mark.parametrize('trino_connection', ['memory/test'], indirect=True)
 @pytest.mark.parametrize('schema', [None, 'test'])
 def test_get_table_names(trino_connection, schema):

trino/sqlalchemy/dialect.py

@@ -190,7 +190,8 @@ def _get_columns(self, connection: Connection, table_name: str, schema: str = No
                 "column_name",
                 "data_type",
                 "column_default",
-                UPPER("is_nullable") AS "is_nullable"
+                UPPER("is_nullable") AS "is_nullable",
+                "comment"
             FROM "information_schema"."columns"
             WHERE "table_schema" = :schema
               AND "table_name" = :table
@@ -205,6 +206,7 @@ def _get_columns(self, connection: Connection, table_name: str, schema: str = No
                 type=datatype.parse_sqltype(record.data_type),
                 nullable=record.is_nullable == "YES",
                 default=record.column_default,
+                comment=record.comment,
             )
             columns.append(column)
         return columns