update loader formatting

Author: Yuan325

src/langchain_google_spanner/loader.py

@@ -46,21 +46,16 @@ def _load_row_to_doc(
 ) -> Document:
     page_content = ""
     if format == "text":
-        page_content = " ".join(str(row[c]) for c in content_columns if c in row)
+        page_content = " ".join(str(row[c]) for c in content_columns)
     elif format == "YAML":
-        page_content = "\n".join(
-            f"{c}: {str(row[c])}" for c in content_columns if c in row
-        )
+        page_content = "\n".join(f"{c}: {str(row[c])}" for c in content_columns)
     elif format == "JSON":
         j = {}
         for c in content_columns:
-            if c in row:
-                j[c] = row[c]
+            j[c] = row[c]
         page_content = json.dumps(j)
     elif format == "CSV":
-        page_content = ", ".join(str(row[c]) for c in content_columns if c in row)
-    if not page_content:
-        raise Exception("column page_content doesn't exist.")
+        page_content = ", ".join(str(row[c]) for c in content_columns)
 
     metadata: Dict[str, Any] = {}
     if metadata_json_column in metadata_columns and row.get(metadata_json_column):

tests/test_spanner_loader.py

@@ -255,6 +255,18 @@ def test_loader_custom_format_error(self, client):
                 client,
                 format="NOT_A_FORMAT",
             )
+            docs = loader.load()
+
+    def test_loader_custom_key_error(self, client):
+        query = f"SELECT * FROM {table_name}"
+        with pytest.raises(Exception):
+            SpannerLoader(
+                instance_id,
+                google_database,
+                query,
+                client,
+                content_columns=["NOT_A_COLUMN"],
+            )
 
     def test_loader_custom_json_metadata(self, client):
         database = client.instance(instance_id).database(google_database)