Pebblo 0.1.18 snowflake v2

libs/community/langchain_community/document_loaders/pebblo.py

@@ -1,10 +1,11 @@
 """Pebblo's safe dataloader is a wrapper for document loaders"""
 
+import ast
 import logging
 import os
 import uuid
 from importlib.metadata import version
-from typing import Dict, Iterator, List, Optional
+from typing import Any, Dict, Iterator, List, Optional
 
 from langchain_core.documents import Document
 
@@ -45,6 +46,7 @@ def __init__(
         classifier_url: Optional[str] = None,
         *,
         classifier_location: str = "local",
+        **kwargs: Any,
     ):
         if not name or not isinstance(name, str):
             raise NameError("Must specify a valid name.")
@@ -57,6 +59,7 @@ def __init__(
         self.source_path = get_loader_full_path(self.loader)
         self.docs: List[Document] = []
         self.docs_with_id: List[IndexedDocument] = []
+        self.kwargs = kwargs.get("kwargs")
         loader_name = str(type(self.loader)).split(".")[-1].split("'")[0]
         self.source_type = get_loader_type(loader_name)
         self.source_path_size = get_source_size(self.source_path)
@@ -162,6 +165,41 @@ def lazy_load(self) -> Iterator[Document]:
     def set_discover_sent(cls) -> None:
         cls._discover_sent = True
 
+    def _get_authfield_from_md(self, doc: Document, auth_field_name: str) -> None:
+        """
+        Extracts the AUTH_FIELD from the metadata and adds it to the document.
+
+        Args:
+            doc (Document): Document object.
+            auth_field_name (str): Name of the AUTH_FIELD.
+        """
+        auth_field_list: List[str] = []
+        # Extract the AUTH_FIELD block
+        auth_field_str = doc.metadata.pop(auth_field_name, "")
+        if auth_field_str:
+            # Convert the string representation of the list to an actual list
+            # using ast.literal_eval
+            auth_field_list = ast.literal_eval(auth_field_str)
+            doc.metadata["authorized_identities"] = auth_field_list
+        else:
+            # logger.info("AUTH_FIELD not found")
+            pass
+
+    def _get_sourcefield_from_md(self, doc: Document, source_field_name: str) -> None:
+        """
+        Extracts the SOURCE_FIELD from the metadata and adds it to the document.
+
+        Args:
+            doc (Document): Document object.
+            source_field_name (str): Name of the SOURCE_FIELD.
+        """
+        auth_field_str = doc.metadata.pop(source_field_name, "")
+        if auth_field_str:
+            doc.metadata["full_path"] = auth_field_str
+        else:
+            # logger.info("SOURCE_FIELD not found")
+            pass
+
     def _get_app_details(self) -> App:
         """Fetch app details. Internal method.
 
@@ -258,6 +296,20 @@ def _add_pebblo_specific_metadata(self, classified_docs: dict) -> None:
         """Add Pebblo specific metadata to documents."""
         for doc in self.docs_with_id:
             doc_metadata = doc.metadata
+            if (
+                self.loader.__class__.__name__ == "SnowflakeLoader"
+                and self.kwargs is not None
+            ):
+                if self.kwargs.get("auth_identities") is not None:
+                    # Snowflake Table column name for authorized_identities
+                    # e.g. values [[email protected], [email protected]]
+                    column_name = self.kwargs.get("auth_identities")
+                    self._get_authfield_from_md(doc, column_name)
+                if self.kwargs.get("source") is not None:
+                    # Snowflake Table column name for authorized_identities
+                    # e.g. values [[email protected], [email protected]]
+                    column_name = self.kwargs.get("source")
+                    self._get_sourcefield_from_md(doc, column_name)
             if self.loader.__class__.__name__ == "SharePointLoader":
                 doc_metadata["full_path"] = get_full_path(
                     doc_metadata.get("source", self.source_path)

libs/community/langchain_community/utilities/pebblo.py

@@ -54,6 +54,7 @@
     "NotionDBLoader",
     "GoogleDriveLoader",
     "SharePointLoader",
+    "SnowflakeLoader",
 ]
 
 LOADER_TYPE_MAPPING = {