add column comment to get_columns method

cockroachdb · Sep 8, 2024 · 726b05d · 726b05d
1 parent c4fe341
commit 726b05d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/sqlalchemy_cockroachdb/base.py b/sqlalchemy_cockroachdb/base.py
@@ -182,9 +182,10 @@ def get_columns(self, conn, table_name, schema=None, **kw):
         if not self._is_v191plus:
             # v2.x does not have is_generated or generation_expression
             sql = (
-                "SELECT column_name, data_type, is_nullable::bool, column_default, "
+                "SELECT column_name, data_type, is_nullable::bool, column_default,"
                 "numeric_precision, numeric_scale, character_maximum_length, "
-                "NULL AS is_generated, NULL AS generation_expression, is_hidden::bool "
+                "NULL AS is_generated, NULL AS generation_expression, is_hidden::bool,"
+                "comment "
                 "FROM information_schema.columns "
                 "WHERE table_schema = :table_schema AND table_name = :table_name "
             )
@@ -200,7 +201,7 @@ def get_columns(self, conn, table_name, schema=None, **kw):
                 "numeric_precision, numeric_scale, character_maximum_length, "
                 "CASE is_generated WHEN 'ALWAYS' THEN true WHEN 'NEVER' THEN false "
                 "ELSE is_generated::bool END AS is_generated, "
-                "generation_expression, is_hidden::bool, crdb_sql_type "
+                "generation_expression, is_hidden::bool, crdb_sql_type, comment "
                 "FROM information_schema.columns "
                 "WHERE table_schema = :table_schema AND table_name = :table_name "
             )
@@ -281,6 +282,7 @@ def get_columns(self, conn, table_name, schema=None, **kw):
                 default=default,
                 autoincrement=autoincrement,
                 is_hidden=row.is_hidden,
+                comment=row.comment,
             )
             if computed is not None:
                 column_info["computed"] = computed

diff --git a/test/test_column_reflect.py b/test/test_column_reflect.py
@@ -44,6 +44,7 @@ def test_reflect_hidden_columns(self):
                     "default": "unique_rowid()",
                     "autoincrement": True,
                     "is_hidden": False,
+                    "comment": None,
                 },
                 {
                     "name": "txt",
@@ -52,6 +53,7 @@ def test_reflect_hidden_columns(self):
                     "default": None,
                     "autoincrement": False,
                     "is_hidden": False,
+                    "comment": None,
                 },
             ],
         )
@@ -66,6 +68,7 @@ def test_reflect_hidden_columns(self):
                     "default": None,
                     "autoincrement": False,
                     "is_hidden": False,
+                    "comment": None,
                 },
             ],
         )
@@ -80,6 +83,7 @@ def test_reflect_hidden_columns(self):
                     "default": None,
                     "autoincrement": False,
                     "is_hidden": False,
+                    "comment": None,
                 },
                 {
                     "name": "rowid",
@@ -88,6 +92,7 @@ def test_reflect_hidden_columns(self):
                     "default": "unique_rowid()",
                     "autoincrement": True,
                     "is_hidden": True,
+                    "comment": None,
                 },
             ],
         )