Fixed appending to tables by adding filtering of None rows (#356)

This PR aims to fix #297 and
#346

It adds a utility method to filter rows that have a column containing
None, this will help Crawlers to not throw an error when a column is
None.

it also checks if the column in the class is Nullable or not, if it's
nullable and the value is None, it's ignored

Author: william-conti

src/databricks/labs/ucx/framework/crawlers.py

@@ -41,6 +41,32 @@ def _schema_for(cls, klass):
             fields.append(f"{f.name} {spark_type}{not_null}")
         return ", ".join(fields)
 
+    @classmethod
+    def _filter_none_rows(cls, rows, full_name):
+        if len(rows) == 0:
+            return rows
+
+        results = []
+        nullable_fields = set()
+
+        for field in dataclasses.fields(rows[0]):
+            if field.default is None:
+                nullable_fields.add(field.name)
+
+        for row in rows:
+            if row is None:
+                continue
+            row_contains_none = False
+            for column, value in dataclasses.asdict(row).items():
+                if value is None and column not in nullable_fields:
+                    logger.warning(f"[{full_name}] Field {column} is None, filtering row")
+                    row_contains_none = True
+                    break
+
+            if not row_contains_none:
+                results.append(row)
+        return results
+
 
 class StatementExecutionBackend(SqlBackend):
     def __init__(self, ws: WorkspaceClient, warehouse_id, *, max_records_per_batch: int = 1000):
@@ -60,7 +86,7 @@ def save_table(self, full_name: str, rows: list[any], mode="append"):
         if mode == "overwrite":
             msg = "Overwrite mode is not yet supported"
             raise NotImplementedError(msg)
-
+        rows = self._filter_none_rows(rows, full_name)
         if len(rows) == 0:
             return
 
@@ -115,6 +141,8 @@ def fetch(self, sql) -> Iterator[any]:
         return self._spark.sql(sql).collect()
 
     def save_table(self, full_name: str, rows: list[any], mode: str = "append"):
+        rows = self._filter_none_rows(rows, full_name)
+
         if len(rows) == 0:
             return
         # pyspark deals well with lists of dataclass instances, as long as schema is provided

tests/unit/framework/test_crawlers.py

@@ -20,6 +20,12 @@ class Foo:
     second: bool
 
 
+@dataclass
+class Baz:
+    first: str
+    second: str = None
+
+
 @dataclass
 class Bar:
     first: str
@@ -204,3 +210,39 @@ def test_runtime_backend_save_table(mocker):
             "first STRING NOT NULL, second BOOLEAN NOT NULL",
         )
         rb._spark.createDataFrame().write.saveAsTable.assert_called_with("a.b.c", mode="append")
+
+
+def test_runtime_backend_save_table_with_row_containing_none(mocker):
+    from unittest import mock
+
+    with mock.patch.dict(os.environ, {"DATABRICKS_RUNTIME_VERSION": "14.0"}):
+        pyspark_sql_session = mocker.Mock()
+        sys.modules["pyspark.sql.session"] = pyspark_sql_session
+
+        rb = RuntimeBackend()
+
+        rb.save_table("a.b.c", [Foo("aaa", True), Foo("bbb", False), Foo("ccc", None)])
+
+        rb._spark.createDataFrame.assert_called_with(
+            [Foo(first="aaa", second=True), Foo(first="bbb", second=False)],
+            "first STRING NOT NULL, second BOOLEAN NOT NULL",
+        )
+        rb._spark.createDataFrame().write.saveAsTable.assert_called_with("a.b.c", mode="append")
+
+
+def test_runtime_backend_save_table_with_row_containing_none_with_nullable_class(mocker):
+    from unittest import mock
+
+    with mock.patch.dict(os.environ, {"DATABRICKS_RUNTIME_VERSION": "14.0"}):
+        pyspark_sql_session = mocker.Mock()
+        sys.modules["pyspark.sql.session"] = pyspark_sql_session
+
+        rb = RuntimeBackend()
+
+        rb.save_table("a.b.c", [Baz("aaa", "ccc"), Baz("bbb", None)])
+
+        rb._spark.createDataFrame.assert_called_with(
+            [Baz(first="aaa", second="ccc"), Baz(first="bbb", second=None)],
+            "first STRING NOT NULL, second STRING",
+        )
+        rb._spark.createDataFrame().write.saveAsTable.assert_called_with("a.b.c", mode="append")