fix: Return None for missing column and partition key comment (#2449)

* Return None for missing column and partition key comment

* fix return type

* extend returned example

---------

Co-authored-by: Robert Schmidtke <[email protected]>
Co-authored-by: kukushking <[email protected]>
Co-authored-by: jaidisido <[email protected]>

Author: 4 people

awswrangler/catalog/_get.py

@@ -961,7 +961,7 @@ def get_columns_comments(
     query_as_of_time: Optional[str] = None,
     catalog_id: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
-) -> Dict[str, str]:
+) -> Dict[str, Optional[str]]:
     """Get all columns comments.
 
     Note
@@ -987,8 +987,8 @@ def get_columns_comments(
 
     Returns
     -------
-    Dict[str, str]
-        Columns comments. e.g. {"col1": "foo boo bar"}.
+    Dict[str, Optional[str]]
+        Columns comments. e.g. {"col1": "foo boo bar", "col2": None}.
 
     Examples
     --------
@@ -1005,12 +1005,12 @@ def get_columns_comments(
             ),
         )
     )
-    comments: Dict[str, str] = {}
+    comments: Dict[str, Optional[str]] = {}
     for c in response["Table"]["StorageDescriptor"]["Columns"]:
-        comments[c["Name"]] = c["Comment"]
+        comments[c["Name"]] = c.get("Comment")
     if "PartitionKeys" in response["Table"]:
         for p in response["Table"]["PartitionKeys"]:
-            comments[p["Name"]] = p["Comment"]
+            comments[p["Name"]] = p.get("Comment")
     return comments