update loader formatting

Author: Yuan325

src/langchain_google_spanner/loader.py

@@ -46,28 +46,23 @@ 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):
         metadata = row[metadata_json_column]
 
     for c in metadata_columns:
-        if c in row and c != metadata_json_column:
+        if c != metadata_json_column:
             metadata[c] = row[c]
 
     return Document(page_content=page_content, metadata=metadata)
@@ -273,9 +268,9 @@ def __init__(
                 "Table doesn't exist. Create table with SpannerDocumentSaver.init_document_table function."
             )
         self._table_fields = [content_column]
-        self._table_fields.append(
-            [n.name for n in table.schema if n.name != metadata_json_column]
-        )
+        for n in table.schema:
+            if n.name != metadata_json_column:
+                self._table_fields.append(n.name)
         self._table_fields.append(metadata_json_column)
 
     def add_documents(self, documents: List[Document]):

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)