You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "seungoh-lee (via GitHub)" <gi...@apache.org> on 2026/04/28 10:24:30 UTC
[PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
seungoh-lee opened a new pull request, #66013:
URL: https://github.com/apache/airflow/pull/66013
Introduces a new transfer operator under the Google provider that exports MongoDB documents to Google Cloud Storage in JSON, CSV, or Parquet format. The operator extends BaseSQLToGCSOperator and accepts either a find filter or an aggregation pipeline through mongo_query, with optional projection and allowDiskUse controls.
A small cursor adapter wraps the pymongo cursor as a DB-API style cursor so the parent operator chunking, schema inference, and upload flow can be reused. BSON-specific values (ObjectId, Decimal128, bytes) are converted to BigQuery-friendly types.
Includes provider.yaml registration, mongo provider as a transitive extra, unit tests, system test example, and how-to documentation.
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [X] Yes — Claude Code (Opus 4.7)
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "seungoh-lee (via GitHub)" <gi...@apache.org>.
seungoh-lee commented on code in PR #66013:
URL: https://github.com/apache/airflow/pull/66013#discussion_r3325237222
##########
providers/google/src/airflow/providers/google/cloud/transfers/mongo_to_gcs.py:
##########
@@ -0,0 +1,204 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, 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.
+"""MongoDB to GCS operator."""
+
+from __future__ import annotations
+
+import base64
+import json
+from collections.abc import Iterator, Sequence
+from datetime import date, datetime, time
+from decimal import Decimal
+from functools import cached_property
+from typing import Any
+
+from bson import ObjectId
+from bson.decimal128 import Decimal128
+
+from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
+from airflow.providers.mongo.hooks.mongo import MongoHook
+
+
+class _MongoCursorAdapter:
+ """
+ Wrap a pymongo cursor as a DB-API 2.0 style cursor.
+
+ ``BaseSQLToGCSOperator`` consumes ``cursor.description`` to derive the
+ BigQuery schema and iterates the cursor expecting tuples. MongoDB documents
+ are dict-shaped and have no fixed schema; this adapter:
+
+ * Peeks the first document to derive ``description`` (column names and
+ Python types).
+ * Re-yields the first document, then iterates the rest.
+ * Converts each document to a tuple in the derived column order, filling
+ missing fields with ``None``.
+ """
+
+ def __init__(self, cursor: Any) -> None:
+ self._cursor = iter(cursor)
+ self._first: dict | None = None
+ self._description: list[tuple] = []
+ try:
+ self._first = next(self._cursor)
+ except StopIteration:
+ return
+ self._description = [
+ (name, type(value), None, None, None, None, True) for name, value in self._first.items()
+ ]
+
+ @property
+ def description(self) -> list[tuple]:
+ return self._description
+
+ def __iter__(self) -> Iterator[tuple]:
+ if self._first is None:
+ return
+ names = [d[0] for d in self._description]
+ yield tuple(self._first.get(n) for n in names)
+ for doc in self._cursor:
+ yield tuple(doc.get(n) for n in names)
+
+
+class MongoToGCSOperator(BaseSQLToGCSOperator):
+ """
+ Copy data from MongoDB to Google Cloud Storage in JSON, CSV or Parquet format.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:MongoToGCSOperator`
+
+ :param mongo_conn_id: Reference to a specific
+ :ref:`Mongo connection <howto/connection:mongo>`.
+ :param mongo_db: The MongoDB database name.
+ :param mongo_collection: The MongoDB collection name.
+ :param mongo_query: A MongoDB find filter (``dict``) or aggregation pipeline
+ (``list``). Defaults to ``{}`` (match all).
+ :param mongo_projection: Optional projection passed to ``find()``. Ignored
+ when ``mongo_query`` is an aggregation pipeline. Accepts a dict
+ (``{"field": 1}``) or list of field names.
+ :param allow_disk_use: Whether to pass ``allowDiskUse=True`` to
+ ``aggregate()``. Defaults to True.
+ """
+
+ ui_color = "#a0e08c"
+
+ template_fields: Sequence[str] = (
+ *BaseSQLToGCSOperator.template_fields,
Review Comment:
dropped the inherited-but-unused sql field by enumerating template_fields
explicitly. Also overrode template_ext (.sql) and template_fields_renderers, since they're equally
meaningless here — this operator is driven by mongo_query, not SQL.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "seungoh-lee (via GitHub)" <gi...@apache.org>.
seungoh-lee commented on PR #66013:
URL: https://github.com/apache/airflow/pull/66013#issuecomment-4394441798
I don't think these test failures are related to my changes. could you please take a look?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #66013:
URL: https://github.com/apache/airflow/pull/66013#issuecomment-4334364476
Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our [Contributors' Guide](https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:
- Pay attention to the quality of your code (ruff, mypy and type annotations). Our [prek-hooks]( https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst#prerequisites-for-prek-hooks) will help you with that.
- In case of a new feature add useful documentation (in docstrings or in `docs/` directory). Adding a new operator? Check this short [guide](https://github.com/apache/airflow/blob/main/airflow-core/docs/howto/custom-operator.rst) Consider adding an example DAG that shows how users should use it.
- Consider using [Breeze environment](https://github.com/apache/airflow/blob/main/dev/breeze/doc/README.rst) for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
- Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
- Please follow [ASF Code of Conduct](https://www.apache.org/foundation/policies/conduct) for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
- Be sure to read the [Airflow Coding style]( https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#coding-style-and-best-practices).
- Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: dev@airflow.apache.org
Slack: https://s.apache.org/airflow-slack
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #66013:
URL: https://github.com/apache/airflow/pull/66013#issuecomment-4529556851
@seungoh-lee A few things need addressing before review — see our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
- :x: **CI fails**: `CI image checks / Build documentation (--spellcheck-only)` (and possibly other checks — see the [Checks tab](https://github.com/apache/airflow/pull/66013/checks) for the full list).
No rush.
---
_Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this [two-stage triage process](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md#why-the-first-pass-is-automated) so that our maintainers' limited time is spent where it matters most: the conversation with you._
---
Drafted-by: Claude Code (Opus 4.7); reviewed by @potiuk before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #66013:
URL: https://github.com/apache/airflow/pull/66013#issuecomment-4394526962
Yeah they don't. Pls rebase with main
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "seungoh-lee (via GitHub)" <gi...@apache.org>.
seungoh-lee commented on PR #66013:
URL: https://github.com/apache/airflow/pull/66013#issuecomment-4562669320
> @seungoh-lee A few things need addressing before review — see our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
>
> * ❌ **CI fails**: `CI image checks / Build documentation (--spellcheck-only)` (and possibly other checks — see the [Checks tab](https://github.com/apache/airflow/pull/66013/checks) for the full list).
>
> No rush.
>
> _Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this [two-stage triage process](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md#why-the-first-pass-is-automated) so that our maintainers' limited time is spent where it matters most: the conversation with you._
>
> Drafted-by: Claude Code (Opus 4.7); reviewed by @potiuk before posting
I'm done reflecting what you said
Can you check it one more time?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #66013:
URL: https://github.com/apache/airflow/pull/66013#issuecomment-4659032985
Awesome work, congrats on your first merged pull request! You are invited to check our [Issue Tracker](https://github.com/apache/airflow/issues) for additional contributions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "seungoh-lee (via GitHub)" <gi...@apache.org>.
seungoh-lee commented on PR #66013:
URL: https://github.com/apache/airflow/pull/66013#issuecomment-4394666082
Sorry, rebased onto latest main
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "seungoh-lee (via GitHub)" <gi...@apache.org>.
seungoh-lee commented on code in PR #66013:
URL: https://github.com/apache/airflow/pull/66013#discussion_r3325233899
##########
providers/google/tests/unit/google/cloud/transfers/test_mongo_to_gcs.py:
##########
@@ -0,0 +1,273 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, 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.
+from __future__ import annotations
+
+import datetime
+from decimal import Decimal
+from unittest import mock
+
+import pytest
+from bson import ObjectId
+from bson.decimal128 import Decimal128
+
+from airflow.providers.google.cloud.transfers.mongo_to_gcs import (
+ MongoToGCSOperator,
+ _MongoCursorAdapter,
+)
+
+TASK_ID = "test-mongo-to-gcs"
+MONGO_CONN_ID = "mongo_test"
+MONGO_DB = "test_db"
+MONGO_COLLECTION = "test_collection"
+BUCKET = "gs://test"
+JSON_FILENAME = "test_{}.ndjson"
+
+
+class TestMongoCursorAdapter:
+ def test_empty_cursor(self):
+ adapter = _MongoCursorAdapter(iter([]))
+ assert adapter.description == []
+ assert list(adapter) == []
+
+ def test_description_derived_from_first_doc(self):
+ docs = [{"name": "alice", "age": 30}, {"name": "bob", "age": 25}]
+ adapter = _MongoCursorAdapter(iter(docs))
+
+ names = [d[0] for d in adapter.description]
+ types = [d[1] for d in adapter.description]
+ assert names == ["name", "age"]
+ assert types == [str, int]
+
+ def test_iteration_returns_tuples_in_description_order(self):
+ docs = [{"name": "alice", "age": 30}, {"name": "bob", "age": 25}]
+ adapter = _MongoCursorAdapter(iter(docs))
+
+ rows = list(adapter)
+ assert rows == [("alice", 30), ("bob", 25)]
+
+ def test_missing_field_filled_with_none(self):
+ docs = [{"name": "alice", "age": 30}, {"name": "bob"}]
+ adapter = _MongoCursorAdapter(iter(docs))
+
+ rows = list(adapter)
+ assert rows == [("alice", 30), ("bob", None)]
+
+
+class TestMongoToGCSOperator:
+ def test_init(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_conn_id=MONGO_CONN_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.task_id == TASK_ID
+ assert op.mongo_conn_id == MONGO_CONN_ID
+ assert op.mongo_db == MONGO_DB
+ assert op.mongo_collection == MONGO_COLLECTION
+ assert op.mongo_query == {}
+ assert op.mongo_projection is None
+ assert op.allow_disk_use is True
+ assert op.bucket == BUCKET
+ assert op.filename == JSON_FILENAME
+
+ def test_field_to_bigquery_known_type(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.field_to_bigquery(("col", int, None, None, None, None, True)) == {
+ "name": "col",
+ "type": "INTEGER",
+ "mode": "NULLABLE",
+ }
+
+ def test_field_to_bigquery_unknown_type_defaults_to_string(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+
+ class _Custom:
+ pass
+
+ assert op.field_to_bigquery(("col", _Custom, None, None, None, None, True))["type"] == "STRING"
+
+ @pytest.mark.parametrize(
+ ("value", "schema_type", "expected"),
+ [
+ (None, None, None),
+ ("text", None, "text"),
+ (42, None, 42),
+ (3.14, None, 3.14),
+ (True, None, True),
+ (Decimal("1.5"), None, 1.5),
+ (datetime.datetime(2023, 1, 2, 3, 4, 5), None, "2023-01-02 03:04:05"),
+ (datetime.date(2023, 1, 2), "DATE", "2023-01-02"),
+ (datetime.date(2023, 1, 2), None, "2023-01-02 00:00:00"),
+ ([1, 2, 3], None, "[1, 2, 3]"),
+ ({"a": 1}, None, '{"a": 1}'),
+ ((1, 2), None, "[1, 2]"),
+ ],
+ )
+ def test_convert_type_basic(self, value, schema_type, expected):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.convert_type(value, schema_type) == expected
+
+ def test_convert_type_object_id(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ oid = ObjectId("507f1f77bcf86cd799439011")
+ assert op.convert_type(oid, None) == "507f1f77bcf86cd799439011"
+
+ def test_convert_type_decimal128(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.convert_type(Decimal128("3.14"), None) == 3.14
+
+ def test_convert_type_bytes_default_base64(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.convert_type(b"hello", None) == "aGVsbG8="
+
+ def test_convert_type_bytes_as_integer(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.convert_type(b"\x00\x01", "INTEGER") == 1
+
+ @mock.patch("airflow.providers.google.cloud.transfers.mongo_to_gcs.MongoHook")
Review Comment:
added autospec=True to the MongoHook/GCSHook patches in the unit tests
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 commented on code in PR #66013:
URL: https://github.com/apache/airflow/pull/66013#discussion_r3324987033
##########
providers/google/src/airflow/providers/google/cloud/transfers/mongo_to_gcs.py:
##########
@@ -0,0 +1,204 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, 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.
+"""MongoDB to GCS operator."""
+
+from __future__ import annotations
+
+import base64
+import json
+from collections.abc import Iterator, Sequence
+from datetime import date, datetime, time
+from decimal import Decimal
+from functools import cached_property
+from typing import Any
+
+from bson import ObjectId
+from bson.decimal128 import Decimal128
+
+from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
+from airflow.providers.mongo.hooks.mongo import MongoHook
+
+
+class _MongoCursorAdapter:
+ """
+ Wrap a pymongo cursor as a DB-API 2.0 style cursor.
+
+ ``BaseSQLToGCSOperator`` consumes ``cursor.description`` to derive the
+ BigQuery schema and iterates the cursor expecting tuples. MongoDB documents
+ are dict-shaped and have no fixed schema; this adapter:
+
+ * Peeks the first document to derive ``description`` (column names and
+ Python types).
+ * Re-yields the first document, then iterates the rest.
+ * Converts each document to a tuple in the derived column order, filling
+ missing fields with ``None``.
+ """
+
+ def __init__(self, cursor: Any) -> None:
+ self._cursor = iter(cursor)
+ self._first: dict | None = None
+ self._description: list[tuple] = []
+ try:
+ self._first = next(self._cursor)
+ except StopIteration:
+ return
+ self._description = [
+ (name, type(value), None, None, None, None, True) for name, value in self._first.items()
+ ]
+
+ @property
+ def description(self) -> list[tuple]:
+ return self._description
+
+ def __iter__(self) -> Iterator[tuple]:
+ if self._first is None:
+ return
+ names = [d[0] for d in self._description]
+ yield tuple(self._first.get(n) for n in names)
+ for doc in self._cursor:
+ yield tuple(doc.get(n) for n in names)
+
+
+class MongoToGCSOperator(BaseSQLToGCSOperator):
Review Comment:
The inheritence looks a bit unnatural, as Mongo is NoSQL.
Please add a clarification about why this is needed in the docstring (maybe in the future we could utilize `BaseNoSQLToGCSOperator` which could be more dedicated).
##########
providers/google/tests/unit/google/cloud/transfers/test_mongo_to_gcs.py:
##########
@@ -0,0 +1,273 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, 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.
+from __future__ import annotations
+
+import datetime
+from decimal import Decimal
+from unittest import mock
+
+import pytest
+from bson import ObjectId
+from bson.decimal128 import Decimal128
+
+from airflow.providers.google.cloud.transfers.mongo_to_gcs import (
+ MongoToGCSOperator,
+ _MongoCursorAdapter,
+)
+
+TASK_ID = "test-mongo-to-gcs"
+MONGO_CONN_ID = "mongo_test"
+MONGO_DB = "test_db"
+MONGO_COLLECTION = "test_collection"
+BUCKET = "gs://test"
+JSON_FILENAME = "test_{}.ndjson"
+
+
+class TestMongoCursorAdapter:
+ def test_empty_cursor(self):
+ adapter = _MongoCursorAdapter(iter([]))
+ assert adapter.description == []
+ assert list(adapter) == []
+
+ def test_description_derived_from_first_doc(self):
+ docs = [{"name": "alice", "age": 30}, {"name": "bob", "age": 25}]
+ adapter = _MongoCursorAdapter(iter(docs))
+
+ names = [d[0] for d in adapter.description]
+ types = [d[1] for d in adapter.description]
+ assert names == ["name", "age"]
+ assert types == [str, int]
+
+ def test_iteration_returns_tuples_in_description_order(self):
+ docs = [{"name": "alice", "age": 30}, {"name": "bob", "age": 25}]
+ adapter = _MongoCursorAdapter(iter(docs))
+
+ rows = list(adapter)
+ assert rows == [("alice", 30), ("bob", 25)]
+
+ def test_missing_field_filled_with_none(self):
+ docs = [{"name": "alice", "age": 30}, {"name": "bob"}]
+ adapter = _MongoCursorAdapter(iter(docs))
+
+ rows = list(adapter)
+ assert rows == [("alice", 30), ("bob", None)]
+
+
+class TestMongoToGCSOperator:
+ def test_init(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_conn_id=MONGO_CONN_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.task_id == TASK_ID
+ assert op.mongo_conn_id == MONGO_CONN_ID
+ assert op.mongo_db == MONGO_DB
+ assert op.mongo_collection == MONGO_COLLECTION
+ assert op.mongo_query == {}
+ assert op.mongo_projection is None
+ assert op.allow_disk_use is True
+ assert op.bucket == BUCKET
+ assert op.filename == JSON_FILENAME
+
+ def test_field_to_bigquery_known_type(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.field_to_bigquery(("col", int, None, None, None, None, True)) == {
+ "name": "col",
+ "type": "INTEGER",
+ "mode": "NULLABLE",
+ }
+
+ def test_field_to_bigquery_unknown_type_defaults_to_string(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+
+ class _Custom:
+ pass
+
+ assert op.field_to_bigquery(("col", _Custom, None, None, None, None, True))["type"] == "STRING"
+
+ @pytest.mark.parametrize(
+ ("value", "schema_type", "expected"),
+ [
+ (None, None, None),
+ ("text", None, "text"),
+ (42, None, 42),
+ (3.14, None, 3.14),
+ (True, None, True),
+ (Decimal("1.5"), None, 1.5),
+ (datetime.datetime(2023, 1, 2, 3, 4, 5), None, "2023-01-02 03:04:05"),
+ (datetime.date(2023, 1, 2), "DATE", "2023-01-02"),
+ (datetime.date(2023, 1, 2), None, "2023-01-02 00:00:00"),
+ ([1, 2, 3], None, "[1, 2, 3]"),
+ ({"a": 1}, None, '{"a": 1}'),
+ ((1, 2), None, "[1, 2]"),
+ ],
+ )
+ def test_convert_type_basic(self, value, schema_type, expected):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.convert_type(value, schema_type) == expected
+
+ def test_convert_type_object_id(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ oid = ObjectId("507f1f77bcf86cd799439011")
+ assert op.convert_type(oid, None) == "507f1f77bcf86cd799439011"
+
+ def test_convert_type_decimal128(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.convert_type(Decimal128("3.14"), None) == 3.14
+
+ def test_convert_type_bytes_default_base64(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.convert_type(b"hello", None) == "aGVsbG8="
+
+ def test_convert_type_bytes_as_integer(self):
+ op = MongoToGCSOperator(
+ task_id=TASK_ID,
+ mongo_db=MONGO_DB,
+ mongo_collection=MONGO_COLLECTION,
+ bucket=BUCKET,
+ filename=JSON_FILENAME,
+ )
+ assert op.convert_type(b"\x00\x01", "INTEGER") == 1
+
+ @mock.patch("airflow.providers.google.cloud.transfers.mongo_to_gcs.MongoHook")
Review Comment:
nit: when patching, it's better to use the `autospec`
##########
providers/google/src/airflow/providers/google/cloud/transfers/mongo_to_gcs.py:
##########
@@ -0,0 +1,204 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, 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.
+"""MongoDB to GCS operator."""
+
+from __future__ import annotations
+
+import base64
+import json
+from collections.abc import Iterator, Sequence
+from datetime import date, datetime, time
+from decimal import Decimal
+from functools import cached_property
+from typing import Any
+
+from bson import ObjectId
+from bson.decimal128 import Decimal128
+
+from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
+from airflow.providers.mongo.hooks.mongo import MongoHook
+
+
+class _MongoCursorAdapter:
+ """
+ Wrap a pymongo cursor as a DB-API 2.0 style cursor.
+
+ ``BaseSQLToGCSOperator`` consumes ``cursor.description`` to derive the
+ BigQuery schema and iterates the cursor expecting tuples. MongoDB documents
+ are dict-shaped and have no fixed schema; this adapter:
+
+ * Peeks the first document to derive ``description`` (column names and
+ Python types).
+ * Re-yields the first document, then iterates the rest.
+ * Converts each document to a tuple in the derived column order, filling
+ missing fields with ``None``.
+ """
+
+ def __init__(self, cursor: Any) -> None:
+ self._cursor = iter(cursor)
+ self._first: dict | None = None
+ self._description: list[tuple] = []
+ try:
+ self._first = next(self._cursor)
+ except StopIteration:
+ return
+ self._description = [
+ (name, type(value), None, None, None, None, True) for name, value in self._first.items()
+ ]
+
+ @property
+ def description(self) -> list[tuple]:
+ return self._description
+
+ def __iter__(self) -> Iterator[tuple]:
+ if self._first is None:
+ return
+ names = [d[0] for d in self._description]
+ yield tuple(self._first.get(n) for n in names)
+ for doc in self._cursor:
+ yield tuple(doc.get(n) for n in names)
+
+
+class MongoToGCSOperator(BaseSQLToGCSOperator):
+ """
+ Copy data from MongoDB to Google Cloud Storage in JSON, CSV or Parquet format.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:MongoToGCSOperator`
+
+ :param mongo_conn_id: Reference to a specific
+ :ref:`Mongo connection <howto/connection:mongo>`.
+ :param mongo_db: The MongoDB database name.
+ :param mongo_collection: The MongoDB collection name.
+ :param mongo_query: A MongoDB find filter (``dict``) or aggregation pipeline
+ (``list``). Defaults to ``{}`` (match all).
+ :param mongo_projection: Optional projection passed to ``find()``. Ignored
+ when ``mongo_query`` is an aggregation pipeline. Accepts a dict
+ (``{"field": 1}``) or list of field names.
+ :param allow_disk_use: Whether to pass ``allowDiskUse=True`` to
+ ``aggregate()``. Defaults to True.
+ """
+
+ ui_color = "#a0e08c"
+
+ template_fields: Sequence[str] = (
+ *BaseSQLToGCSOperator.template_fields,
Review Comment:
`sql` is a template_field of the parent, so I would exclude it here
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "seungoh-lee (via GitHub)" <gi...@apache.org>.
seungoh-lee commented on code in PR #66013:
URL: https://github.com/apache/airflow/pull/66013#discussion_r3325241478
##########
providers/google/src/airflow/providers/google/cloud/transfers/mongo_to_gcs.py:
##########
@@ -0,0 +1,204 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, 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.
+"""MongoDB to GCS operator."""
+
+from __future__ import annotations
+
+import base64
+import json
+from collections.abc import Iterator, Sequence
+from datetime import date, datetime, time
+from decimal import Decimal
+from functools import cached_property
+from typing import Any
+
+from bson import ObjectId
+from bson.decimal128 import Decimal128
+
+from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
+from airflow.providers.mongo.hooks.mongo import MongoHook
+
+
+class _MongoCursorAdapter:
+ """
+ Wrap a pymongo cursor as a DB-API 2.0 style cursor.
+
+ ``BaseSQLToGCSOperator`` consumes ``cursor.description`` to derive the
+ BigQuery schema and iterates the cursor expecting tuples. MongoDB documents
+ are dict-shaped and have no fixed schema; this adapter:
+
+ * Peeks the first document to derive ``description`` (column names and
+ Python types).
+ * Re-yields the first document, then iterates the rest.
+ * Converts each document to a tuple in the derived column order, filling
+ missing fields with ``None``.
+ """
+
+ def __init__(self, cursor: Any) -> None:
+ self._cursor = iter(cursor)
+ self._first: dict | None = None
+ self._description: list[tuple] = []
+ try:
+ self._first = next(self._cursor)
+ except StopIteration:
+ return
+ self._description = [
+ (name, type(value), None, None, None, None, True) for name, value in self._first.items()
+ ]
+
+ @property
+ def description(self) -> list[tuple]:
+ return self._description
+
+ def __iter__(self) -> Iterator[tuple]:
+ if self._first is None:
+ return
+ names = [d[0] for d in self._description]
+ yield tuple(self._first.get(n) for n in names)
+ for doc in self._cursor:
+ yield tuple(doc.get(n) for n in names)
+
+
+class MongoToGCSOperator(BaseSQLToGCSOperator):
Review Comment:
added a note to the class docstring explaining that
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add MongoToGCSOperator to copy MongoDB collections to GCS [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk merged PR #66013:
URL: https://github.com/apache/airflow/pull/66013
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org