Move DTE parsing into _parse.py

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse Whitehouse

src/databricks/sqlalchemy/_parse.py

@@ -251,6 +251,20 @@ def match_dte_rows_by_value(dte_output: List[Dict[str, str]], match: str) -> Lis
     return output_rows
 
 
+def match_dte_rows_by_key(dte_output: List[Dict[str, str]], match: str) -> List[dict]:
+    """Return a list of dictionaries containing only the col_name:data_type pairs where the `col_name`
+    value contains the match argument.
+    """
+
+    output_rows = []
+
+    for row_dict in dte_output:
+        if match in row_dict["col_name"]:
+            output_rows.append(row_dict)
+
+    return output_rows
+
+
 def get_fk_strings_from_dte_output(dte_output: List[Dict[str, str]]) -> List[dict]:
     """If the DESCRIBE TABLE EXTENDED output contains foreign key constraints, return a list of dictionaries,
     one dictionary per defined constraint
@@ -275,6 +289,15 @@ def get_pk_strings_from_dte_output(
     return output
 
 
+def get_comment_from_dte_output(dte_output: List[Dict[str, str]]) -> Optional[str]:
+    """Returns the value of the first "Comment" col_name data in dte_output"""
+    output = match_dte_rows_by_key(dte_output, "Comment")
+    if not output:
+        return None
+    else:
+        return output[0]["data_type"]
+
+
 # The keys of this dictionary are the values we expect to see in a
 # TGetColumnsRequest's .TYPE_NAME attribute.
 # These are enumerated in ttypes.py as class TTypeId.

src/databricks/sqlalchemy/base.py

@@ -10,6 +10,7 @@
     build_pk_dict,
     get_fk_strings_from_dte_output,
     get_pk_strings_from_dte_output,
+    get_comment_from_dte_output,
     parse_column_info_from_tgetcolumnsresponse,
 )
 
@@ -387,16 +388,16 @@ def get_table_comment(
             table_name=table_name,
             schema_name=schema,
         )
-        # Type ignore is because mypy knows that self._describe_table_extended *can*
-        # return None (even though it never will since expect_result defaults to True)
-        comment_row: Dict[str, str] = next(
-            filter(lambda r: r["col_name"] == "Comment", result), None
-        )  # type: ignore
-        return (
-            {"text": comment_row["data_type"]}
-            if comment_row
-            else ReflectionDefaults.table_comment()
-        )
+
+        if result is None:
+            return ReflectionDefaults.table_comment()
+
+        comment = get_comment_from_dte_output(result)
+
+        if comment:
+            return dict(text=comment)
+        else:
+            return ReflectionDefaults.table_comment()
 
 
 @event.listens_for(Engine, "do_connect")