added create_from to pipeline source

Author: daniel-sanche

google/cloud/firestore_v1/pipeline_source.py

@@ -22,6 +22,9 @@
     from google.cloud.firestore_v1.client import Client
     from google.cloud.firestore_v1.async_client import AsyncClient
     from google.cloud.firestore_v1.base_document import BaseDocumentReference
+    from google.cloud.firestore_v1.base_query import BaseQuery
+    from google.cloud.firestore_v1.base_aggregation import BaseAggregationQuery
+    from google.cloud.firestore_v1.base_collection import BaseCollectionReference
 
 
 PipelineType = TypeVar("PipelineType", bound=_BasePipeline)
@@ -43,6 +46,21 @@ def __init__(self, client: Client | AsyncClient):
     def _create_pipeline(self, source_stage):
         return self.client._pipeline_cls._create_with_stages(self.client, source_stage)
 
+    def create_from(self, query: "BaseQuery" | "BaseAggregationQuery" | "BaseCollectionReference") -> PipelineType:
+        """
+        Create a pipeline from an existing query
+
+        Queries containing a `cursor` or `limit_to_last` are not currently supported
+
+        Args:
+            query: the query to build the pipeline off of
+        Raises:
+            - NotImplementedError: raised if the query contains a `cursor` or `limit_to_last`
+        Returns:
+            a new pipeline instance representing the query
+        """
+        return query._build_pipeline(self)
+
     def collection(self, path: str | tuple[str]) -> PipelineType:
         """
         Creates a new Pipeline that operates on a specified Firestore collection.

tests/unit/v1/test_pipeline_source.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License
+import mock
 
 from google.cloud.firestore_v1.pipeline_source import PipelineSource
 from google.cloud.firestore_v1.pipeline import Pipeline
@@ -19,6 +20,8 @@
 from google.cloud.firestore_v1.async_client import AsyncClient
 from google.cloud.firestore_v1 import pipeline_stages as stages
 from google.cloud.firestore_v1.base_document import BaseDocumentReference
+from google.cloud.firestore_v1.query import Query
+from google.cloud.firestore_v1.async_query import AsyncQuery
 
 
 class TestPipelineSource:
@@ -27,6 +30,9 @@ class TestPipelineSource:
     def _make_client(self):
         return Client()
 
+    def _make_query(self):
+        return Query(mock.Mock())
+
     def test_make_from_client(self):
         instance = self._make_client().pipeline()
         assert isinstance(instance, PipelineSource)
@@ -36,6 +42,23 @@ def test_create_pipeline(self):
         ppl = instance._create_pipeline(None)
         assert isinstance(ppl, self._expected_pipeline_type)
 
+    def test_create_from_mock(self):
+        mock_query = mock.Mock()
+        expected = object()
+        mock_query._build_pipeline.return_value = expected
+        instance = self._make_client().pipeline()
+        got = instance.create_from(mock_query)
+        assert got == expected
+        assert mock_query._build_pipeline.call_count == 1
+        assert mock_query._build_pipeline.call_args_list[0][0][0] == instance
+
+    def test_create_from_query(self):
+        query = self._make_query()
+        instance = self._make_client().pipeline()
+        ppl = instance.create_from(query)
+        assert isinstance(ppl, self._expected_pipeline_type)
+        assert len(ppl.stages) == 1
+
     def test_collection(self):
         instance = self._make_client().pipeline()
         ppl = instance.collection("path")
@@ -98,3 +121,6 @@ class TestPipelineSourceWithAsyncClient(TestPipelineSource):
 
     def _make_client(self):
         return AsyncClient()
+
+    def _make_query(self):
+        return AsyncQuery(mock.Mock())