fix(supabase): address reviewer feedback - lazy init, DocumentStore base class, filters, async

Author: ShubhamGond105

integrations/supabase/src/haystack_integrations/components/retrievers/supabase/groonga_retriever.py

@@ -15,27 +15,31 @@
 @component
 class SupabaseGroongaRetriever:
     """
-        Retrieves documents from SupabaseGroongaDocumentStore using PGroonga full-text search.
+    Retrieves documents from SupabaseGroongaDocumentStore using PGroonga full-text search.
 
-        This retriever works without embeddings — it searches documents using plain text queries.
-        It can be used alongside SupabasePgvectorEmbeddingRetriever in hybrid search pipelines.
+    This retriever works without embeddings — it searches documents using plain text queries.
+    It can be used alongside SupabasePgvectorEmbeddingRetriever in hybrid search pipelines.
 
-        Example usage:
+    Note: async operations are not supported as the supabase-py sync client does not expose
+    awaitable query methods. Use the sync run() method instead.
 
-    ```python
-        from haystack_integrations.document_stores.supabase import SupabaseGroongaDocumentStore
-        from haystack_integrations.components.retrievers.supabase import SupabaseGroongaRetriever
-        from haystack.utils import Secret
-
-        document_store = SupabaseGroongaDocumentStore(
-            supabase_url="https://<project>.supabase.co",
-            supabase_key=Secret.from_env_var("SUPABASE_SERVICE_KEY"),
-            table_name="haystack_fts_documents",
-        )
+    Example usage:
 
-        retriever = SupabaseGroongaRetriever(document_store=document_store, top_k=10)
-        result = retriever.run(query="python programming")
-        print(result["documents"])
+    ```python
+    from haystack_integrations.document_stores.supabase import SupabaseGroongaDocumentStore
+    from haystack_integrations.components.retrievers.supabase import SupabaseGroongaRetriever
+    from haystack.utils import Secret
+
+    document_store = SupabaseGroongaDocumentStore(
+        supabase_url="https://<project>.supabase.co",
+        supabase_key=Secret.from_env_var("SUPABASE_SERVICE_KEY"),
+        table_name="haystack_fts_documents",
+    )
+    document_store.warm_up()
+
+    retriever = SupabaseGroongaRetriever(document_store=document_store, top_k=10)
+    result = retriever.run(query="python programming")
+    print(result["documents"])
     ```
     """
 
@@ -85,15 +89,7 @@ def run(
         if not query:
             return {"documents": []}
 
-        # Handle filter policy
-        if filters is not None:
-            if self.filter_policy == FilterPolicy.MERGE:
-                merged_filters = {**self.filters, **filters}
-            else:
-                merged_filters = filters
-        else:
-            merged_filters = self.filters
-
+        merged_filters = self._merge_filters(filters)
         effective_top_k = top_k if top_k is not None else self.top_k
 
         documents = self.document_store._groonga_retrieval(
@@ -104,6 +100,41 @@ def run(
 
         return {"documents": documents}
 
+    @component.output_types(documents=list[Document])
+    async def run_async(
+        self,
+        query: str,
+        filters: dict[str, Any] | None = None,
+        top_k: int | None = None,
+    ) -> dict[str, list[Document]]:
+        """
+        Async version of run(). 
+
+        Note: supabase-py's sync client does not support native async queries.
+        This method runs the synchronous retrieval and returns the result.
+        For fully async support, consider using acreate_client() from supabase-py
+        and refactoring the document store accordingly.
+
+        :param query: The text query to search for.
+        :param filters: Optional runtime filters. Merged or replaced based on filter_policy.
+        :param top_k: Optional override for maximum number of documents to return.
+        :returns: Dictionary with key "documents" containing list of matching Documents.
+        """
+        return self.run(query=query, filters=filters, top_k=top_k)
+
+    def _merge_filters(self, runtime_filters: dict[str, Any] | None) -> dict[str, Any]:
+        """
+        Merges runtime filters with init filters based on filter_policy.
+
+        :param runtime_filters: Filters passed at runtime.
+        :returns: Merged filters dictionary.
+        """
+        if runtime_filters is not None:
+            if self.filter_policy == FilterPolicy.MERGE:
+                return {**self.filters, **runtime_filters}
+            return runtime_filters
+        return self.filters
+
     def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.
@@ -131,4 +162,4 @@ def from_dict(cls, data: dict[str, Any]) -> "SupabaseGroongaRetriever":
         data["init_parameters"]["document_store"] = SupabaseGroongaDocumentStore.from_dict(doc_store_params)
         if filter_policy := data["init_parameters"].get("filter_policy"):
             data["init_parameters"]["filter_policy"] = FilterPolicy.from_str(filter_policy)
-        return default_from_dict(cls, data)
+        return default_from_dict(cls, data)

integrations/supabase/src/haystack_integrations/document_stores/supabase/groonga_document_store.py

@@ -2,42 +2,43 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any
+from typing import Any, Optional
 
 from haystack import default_from_dict, default_to_dict, logging
 from haystack.dataclasses import Document
 from haystack.document_stores.errors import DuplicateDocumentError
-from haystack.document_stores.types import DuplicatePolicy
+from haystack.document_stores.types import DocumentStore, DuplicatePolicy
 from haystack.utils.auth import Secret, deserialize_secrets_inplace
 
-from supabase import create_client
+from supabase import create_client, Client
 
 logger = logging.getLogger(__name__)
 
 
-class SupabaseGroongaDocumentStore:
+class SupabaseGroongaDocumentStore(DocumentStore):
     """
-        A Document Store for Supabase using PGroonga for full-text search.
+    A Document Store for Supabase using PGroonga for full-text search.
 
-        PGroonga is a PostgreSQL extension for fast, multilingual full-text search.
-        Unlike vector search, this store works with plain text queries — no embeddings needed.
+    PGroonga is a PostgreSQL extension for fast, multilingual full-text search.
+    Unlike vector search, this store works with plain text queries — no embeddings needed.
 
-        Prerequisites:
-        - A Supabase project with PGroonga extension enabled.
-        - Enable PGroonga in your Supabase project by running:
-          `CREATE EXTENSION IF NOT EXISTS pgroonga;`
+    Prerequisites:
+    - A Supabase project with PGroonga extension enabled.
+    - Enable PGroonga in your Supabase project by running:
+      `CREATE EXTENSION IF NOT EXISTS pgroonga;`
 
-        Example usage:
+    Example usage:
 
     ```python
-        from haystack_integrations.document_stores.supabase import SupabaseGroongaDocumentStore
-        from haystack.utils import Secret
-
-        document_store = SupabaseGroongaDocumentStore(
-            supabase_url="https://<project>.supabase.co",
-            supabase_key=Secret.from_env_var("SUPABASE_SERVICE_KEY"),
-            table_name="haystack_fts_documents",
-        )
+    from haystack_integrations.document_stores.supabase import SupabaseGroongaDocumentStore
+    from haystack.utils import Secret
+
+    document_store = SupabaseGroongaDocumentStore(
+        supabase_url="https://<project>.supabase.co",
+        supabase_key=Secret.from_env_var("SUPABASE_SERVICE_KEY"),
+        table_name="haystack_fts_documents",
+    )
+    document_store.warm_up()
     ```
     """
 
@@ -52,6 +53,8 @@ def __init__(
         """
         Creates a new SupabaseGroongaDocumentStore instance.
 
+        Note: Call warm_up() before using the store to initialize the client and table.
+
         :param supabase_url: The URL of your Supabase project.
             Format: `https://<project-ref>.supabase.co`
         :param supabase_key: The service role key for your Supabase project.
@@ -66,11 +69,17 @@ def __init__(
         self.table_name = table_name
         self.recreate_table = recreate_table
 
-        # Connect to Supabase
+        # Client is initialized lazily in warm_up()
+        self._client: Optional[Client] = None
+
+    def warm_up(self) -> None:
+        """
+        Initializes the Supabase client and sets up the table.
+
+        Must be called before using the document store.
+        """
         key = self.supabase_key.resolve_value() or ""
         self._client = create_client(self.supabase_url, key)
-
-        # Set up the table
         self._setup_table()
 
     def _setup_table(self) -> None:
@@ -79,6 +88,8 @@ def _setup_table(self) -> None:
 
         If recreate_table is True, drops and recreates the table.
         """
+        assert self._client is not None, "Call warm_up() before using the document store."
+
         if self.recreate_table:
             self._client.rpc("exec_sql", {"query": f"DROP TABLE IF EXISTS {self.table_name};"}).execute()
 
@@ -103,36 +114,80 @@ def _setup_table(self) -> None:
 
     def count_documents(self) -> int:
         """
-         the number of documents in the store.
+        Returns the number of documents in the store.
 
         :returns: Number of documents.
         """
+        assert self._client is not None, "Call warm_up() before using the document store."
         result = self._client.table(self.table_name).select("*", count="exact").execute()
         return int(result.count) if result.count is not None else 0
 
-    def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Document]:  # noqa: ARG002
+    def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Document]:
         """
         Returns documents matching the given filters.
 
-        :param filters: Optional dictionary of filters.
+        Supported filters: equality filters on `id`, `content`, and `meta` fields.
+
+        :param filters: Optional dictionary of filters. Example: {"field": "meta.language", "operator": "==", "value": "en"}
         :returns: List of matching Document objects.
         """
+        assert self._client is not None, "Call warm_up() before using the document store."
+
         query = self._client.table(self.table_name).select("*")
+
+        if filters:
+            query = self._apply_filters(query, filters)
+
         result = query.execute()
         return [self._to_haystack_document(row) for row in result.data if isinstance(row, dict)]
 
+    def _apply_filters(self, query: Any, filters: dict[str, Any]) -> Any:
+        """
+        Applies filters to a Supabase query.
+
+        :param query: The Supabase query builder.
+        :param filters: Dictionary of filters to apply.
+        :returns: The query with filters applied.
+        """
+        operator = filters.get("operator", "AND")
+        conditions = filters.get("conditions", [])
+
+        for condition in conditions:
+            field = condition.get("field", "")
+            op = condition.get("operator", "==")
+            value = condition.get("value")
+
+            # Handle nested meta fields e.g. "meta.language"
+            if field.startswith("meta."):
+                meta_key = field[len("meta."):]
+                if op == "==":
+                    query = query.eq(f"meta->>'{meta_key}'", value)
+                elif op == "!=":
+                    query = query.neq(f"meta->>'{meta_key}'", value)
+            else:
+                if op == "==":
+                    query = query.eq(field, value)
+                elif op == "!=":
+                    query = query.neq(field, value)
+                elif op == "in":
+                    query = query.in_(field, value)
+
+        return query
+
     def write_documents(
         self,
         documents: list[Document],
-        policy: DuplicatePolicy = DuplicatePolicy.NONE,
+        policy: DuplicatePolicy = DuplicatePolicy.FAIL,
     ) -> int:
         """
         Writes documents to the store.
 
         :param documents: List of Haystack Document objects to write.
-        :param policy: How to handle duplicate documents.
+        :param policy: How to handle duplicate documents. Defaults to DuplicatePolicy.FAIL.
         :returns: Number of documents written.
         """
+        assert self._client is not None, "Call warm_up() before using the document store."
+
         if not documents:
             return 0
 
@@ -171,6 +226,8 @@ def delete_documents(self, document_ids: list[str]) -> None:
 
         :param document_ids: List of document IDs to delete.
         """
+        assert self._client is not None, "Call warm_up() before using the document store."
+
         if not document_ids:
             return
         self._client.table(self.table_name).delete().in_("id", document_ids).execute()
@@ -179,21 +236,68 @@ def _groonga_retrieval(
         self,
         query: str,
         top_k: int = 10,
-        filters: dict[str, Any] | None = None,  # noqa: ARG002
+        filters: dict[str, Any] | None = None,
     ) -> list[Document]:
         """
         Searches documents using PGroonga full-text search.
 
         :param query: The text query to search for.
         :param top_k: Maximum number of results to return.
-        :param filters: Optional filters to apply.
+        :param filters: Optional filters to apply after retrieval.
         :returns: List of matching Document objects ranked by relevance.
         """
+        assert self._client is not None, "Call warm_up() before using the document store."
+
         result = self._client.rpc(
             "groonga_search", {"query_text": query, "table": self.table_name, "top_k": top_k}
         ).execute()
 
-        return [self._to_haystack_document(row) for row in (result.data or []) if isinstance(row, dict)]
+        documents = [self._to_haystack_document(row) for row in (result.data or []) if isinstance(row, dict)]
+
+        # Apply filters post-retrieval if provided
+        if filters:
+            documents = self._filter_documents_in_memory(documents, filters)
+
+        return documents
+
+    def _filter_documents_in_memory(self, documents: list[Document], filters: dict[str, Any]) -> list[Document]:
+        """
+        Filters a list of documents in memory based on the given filters.
+
+        :param documents: List of documents to filter.
+        :param filters: Dictionary of filters to apply.
+        :returns: Filtered list of documents.
+        """
+        conditions = filters.get("conditions", [])
+        filtered = []
+
+        for doc in documents:
+            match = True
+            for condition in conditions:
+                field = condition.get("field", "")
+                op = condition.get("operator", "==")
+                value = condition.get("value")
+
+                if field.startswith("meta."):
+                    meta_key = field[len("meta."):]
+                    doc_value = doc.meta.get(meta_key)
+                else:
+                    doc_value = getattr(doc, field, None)
+
+                if op == "==" and doc_value != value:
+                    match = False
+                    break
+                elif op == "!=" and doc_value == value:
+                    match = False
+                    break
+                elif op == "in" and doc_value not in value:
+                    match = False
+                    break
+
+            if match:
+                filtered.append(doc)
+
+        return filtered
 
     def _to_haystack_document(self, row: dict[str, Any]) -> Document:
         """
@@ -232,4 +336,4 @@ def from_dict(cls, data: dict[str, Any]) -> "SupabaseGroongaDocumentStore":
         :returns: Deserialized component.
         """
         deserialize_secrets_inplace(data["init_parameters"], ["supabase_key"])
-        return default_from_dict(cls, data)
+        return default_from_dict(cls, data)