shahar1 commented on code in PR #74664:
URL: https://github.com/apache/airflow/pull/74664#discussion_r4500476331
##########
providers/google/src/airflow/providers/google/firebase/hooks/firestore.py:
##########
@@ -84,9 +91,188 @@ def get_conn(self):
self._conn = build_from_document(non_authorized_conn._rootDesc,
http=http_authorized)
return self._conn
+ def get_operation(self, operation_name: str) -> dict[str, Any]:
+ """
+ Retrieve the current state of a long-running Firestore operation.
+
+ :param operation_name: The resource name of the operation.
+ """
+ service = self.get_conn()
+ return (
+ service.projects()
+ .databases()
+ .operations()
+ .get(name=operation_name)
+ .execute(num_retries=self.num_retries)
+ )
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def list_operations(
+ self, database_id: str = "(default)", project_id: str =
PROVIDE_PROJECT_ID
+ ) -> list[dict[str, Any]]:
+ """
+ List long-running operations on the specified Firestore database.
+
+ :param database_id: The Database ID.
+ :param project_id: Optional, Google Cloud Project project_id where the
database belongs.
+ If set to None or missing, the default project_id from the Google
Cloud connection is used.
+ """
+ service = self.get_conn()
+ name = f"projects/{project_id}/databases/{database_id}"
+ operations_resource = service.projects().databases().operations()
+ all_operations: list[dict[str, Any]] = []
+ page_token: str | None = None
+
+ while True:
+ list_kwargs: dict[str, Any] = {"name": name}
+ if page_token:
+ list_kwargs["pageToken"] = page_token
+ response =
operations_resource.list(**list_kwargs).execute(num_retries=self.num_retries)
+ if not isinstance(response, dict):
+ break
+ operations = response.get("operations", [])
+ if isinstance(operations, list):
+ all_operations.extend(operations)
+ next_page_token = response.get("nextPageToken")
+ if not isinstance(next_page_token, str) or not next_page_token:
+ break
+ page_token = next_page_token
+
+ return all_operations
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def find_matching_export_operation(
+ self,
+ body: dict[str, Any],
+ database_id: str = "(default)",
+ project_id: str = PROVIDE_PROJECT_ID,
+ include_completed: bool = True,
+ ) -> dict[str, Any] | None:
+ """
+ Find an existing running or completed export operation matching the
export request body.
+
+ :param body: The export request body.
+ :param database_id: The Database ID.
+ :param project_id: Optional, Google Cloud Project project_id where the
database belongs.
+ :param include_completed: Whether to match already completed
operations in addition to
+ in-progress operations.
+ """
+ target_uri = body.get("outputUriPrefix")
+ if not target_uri or not isinstance(target_uri, str):
+ return None
+ target_collections = sorted(body.get("collectionIds") or [])
+ target_namespaces = sorted(body.get("namespaceIds") or [])
+
+ completed_match: dict[str, Any] | None = None
+ for op in self.list_operations(database_id=database_id,
project_id=project_id):
+ if not isinstance(op, dict):
+ continue
+ metadata = op.get("metadata")
+ if not isinstance(metadata, dict):
+ metadata = {}
+
+ meta_type = metadata.get("@type", "")
+ if meta_type and not meta_type.endswith("ExportDocumentsMetadata"):
+ continue
+
+ op_state = metadata.get("operationState", "")
+ if op.get("error") or op_state in {"FAILED", "CANCELLED",
"CANCELLING"}:
+ continue
+
+ op_collections = sorted(metadata.get("collectionIds") or [])
+ if target_collections != op_collections:
+ continue
+
+ op_namespaces = sorted(metadata.get("namespaceIds") or [])
+ if target_namespaces != op_namespaces:
+ continue
+
+ response_dict = op.get("response")
+ op_uri = metadata.get("outputUriPrefix") or (
+ response_dict.get("outputUriPrefix") if
isinstance(response_dict, dict) else None
+ )
+ if not isinstance(op_uri, str) or op_uri.rstrip("/") !=
target_uri.rstrip("/"):
Review Comment:
**Minor.** A bucket-only URI won't match here. The API accepts
`gs://BUCKET_NAME` and documents that "a prefix will be generated based on the
start time", so the operation should report the generated prefix rather than
the requested one. A retry would then start a second full export while the
first is still running. I haven't confirmed this against the live API.
That is no worse than today, but the reattach behavior doesn't apply to this
form. Please state that in the `reattach_on_restart` docs, or log a warning
when `target_uri` has no path.
##########
providers/google/src/airflow/providers/google/firebase/hooks/firestore.py:
##########
@@ -84,9 +91,188 @@ def get_conn(self):
self._conn = build_from_document(non_authorized_conn._rootDesc,
http=http_authorized)
return self._conn
+ def get_operation(self, operation_name: str) -> dict[str, Any]:
+ """
+ Retrieve the current state of a long-running Firestore operation.
+
+ :param operation_name: The resource name of the operation.
+ """
+ service = self.get_conn()
+ return (
+ service.projects()
+ .databases()
+ .operations()
+ .get(name=operation_name)
+ .execute(num_retries=self.num_retries)
+ )
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def list_operations(
+ self, database_id: str = "(default)", project_id: str =
PROVIDE_PROJECT_ID
+ ) -> list[dict[str, Any]]:
+ """
+ List long-running operations on the specified Firestore database.
+
+ :param database_id: The Database ID.
+ :param project_id: Optional, Google Cloud Project project_id where the
database belongs.
+ If set to None or missing, the default project_id from the Google
Cloud connection is used.
+ """
+ service = self.get_conn()
+ name = f"projects/{project_id}/databases/{database_id}"
+ operations_resource = service.projects().databases().operations()
+ all_operations: list[dict[str, Any]] = []
+ page_token: str | None = None
+
+ while True:
+ list_kwargs: dict[str, Any] = {"name": name}
+ if page_token:
+ list_kwargs["pageToken"] = page_token
+ response =
operations_resource.list(**list_kwargs).execute(num_retries=self.num_retries)
+ if not isinstance(response, dict):
Review Comment:
**Nit.** This guard turns a malformed response into "no operations", so the
task would go on to submit a new export. The `isinstance` guards in this method
also appear to be what keeps the existing tests such as
`test_immediately_complete` passing. Those tests mock `get_conn`, so
`list().execute()` returns a `MagicMock`, and without the guards the truthy
page token would keep the `while True` loop running forever.
I'd drop the guards and stub `list` in those tests.
##########
providers/google/src/airflow/providers/google/firebase/hooks/firestore.py:
##########
@@ -84,9 +91,188 @@ def get_conn(self):
self._conn = build_from_document(non_authorized_conn._rootDesc,
http=http_authorized)
return self._conn
+ def get_operation(self, operation_name: str) -> dict[str, Any]:
+ """
+ Retrieve the current state of a long-running Firestore operation.
+
+ :param operation_name: The resource name of the operation.
+ """
+ service = self.get_conn()
+ return (
+ service.projects()
+ .databases()
+ .operations()
+ .get(name=operation_name)
+ .execute(num_retries=self.num_retries)
+ )
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def list_operations(
+ self, database_id: str = "(default)", project_id: str =
PROVIDE_PROJECT_ID
+ ) -> list[dict[str, Any]]:
+ """
+ List long-running operations on the specified Firestore database.
+
+ :param database_id: The Database ID.
+ :param project_id: Optional, Google Cloud Project project_id where the
database belongs.
+ If set to None or missing, the default project_id from the Google
Cloud connection is used.
+ """
+ service = self.get_conn()
+ name = f"projects/{project_id}/databases/{database_id}"
+ operations_resource = service.projects().databases().operations()
+ all_operations: list[dict[str, Any]] = []
+ page_token: str | None = None
+
+ while True:
+ list_kwargs: dict[str, Any] = {"name": name}
+ if page_token:
+ list_kwargs["pageToken"] = page_token
+ response =
operations_resource.list(**list_kwargs).execute(num_retries=self.num_retries)
+ if not isinstance(response, dict):
+ break
+ operations = response.get("operations", [])
+ if isinstance(operations, list):
+ all_operations.extend(operations)
+ next_page_token = response.get("nextPageToken")
+ if not isinstance(next_page_token, str) or not next_page_token:
+ break
+ page_token = next_page_token
+
+ return all_operations
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def find_matching_export_operation(
+ self,
+ body: dict[str, Any],
+ database_id: str = "(default)",
+ project_id: str = PROVIDE_PROJECT_ID,
+ include_completed: bool = True,
+ ) -> dict[str, Any] | None:
+ """
+ Find an existing running or completed export operation matching the
export request body.
+
+ :param body: The export request body.
+ :param database_id: The Database ID.
+ :param project_id: Optional, Google Cloud Project project_id where the
database belongs.
+ :param include_completed: Whether to match already completed
operations in addition to
+ in-progress operations.
+ """
+ target_uri = body.get("outputUriPrefix")
+ if not target_uri or not isinstance(target_uri, str):
+ return None
+ target_collections = sorted(body.get("collectionIds") or [])
+ target_namespaces = sorted(body.get("namespaceIds") or [])
+
+ completed_match: dict[str, Any] | None = None
+ for op in self.list_operations(database_id=database_id,
project_id=project_id):
+ if not isinstance(op, dict):
+ continue
+ metadata = op.get("metadata")
+ if not isinstance(metadata, dict):
+ metadata = {}
+
+ meta_type = metadata.get("@type", "")
+ if meta_type and not meta_type.endswith("ExportDocumentsMetadata"):
+ continue
+
+ op_state = metadata.get("operationState", "")
+ if op.get("error") or op_state in {"FAILED", "CANCELLED",
"CANCELLING"}:
+ continue
+
+ op_collections = sorted(metadata.get("collectionIds") or [])
+ if target_collections != op_collections:
+ continue
+
+ op_namespaces = sorted(metadata.get("namespaceIds") or [])
+ if target_namespaces != op_namespaces:
+ continue
+
+ response_dict = op.get("response")
+ op_uri = metadata.get("outputUriPrefix") or (
+ response_dict.get("outputUriPrefix") if
isinstance(response_dict, dict) else None
+ )
+ if not isinstance(op_uri, str) or op_uri.rstrip("/") !=
target_uri.rstrip("/"):
+ continue
+
+ is_done = bool(op.get("done"))
+ if not is_done:
+ return op
+ if include_completed and completed_match is None:
+ completed_match = op
+
+ return completed_match
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def start_export_documents(
+ self,
+ body: dict[str, Any],
+ database_id: str = "(default)",
+ project_id: str = PROVIDE_PROJECT_ID,
+ reattach_on_restart: bool = True,
+ ) -> dict[str, Any]:
+ """
+ Start a Firestore export or return an existing running/completed
operation when reattaching.
+
+ :param body: The request body.
+ :param database_id: The Database ID.
+ :param project_id: Optional, Google Cloud Project project_id where the
database belongs.
+ :param reattach_on_restart: If True, check for an existing in-progress
export operation
+ matching ``body`` before submitting a new export request, and
recover if the export
+ already completed on a prior attempt.
+ """
+ if reattach_on_restart:
+ existing_op = self.find_matching_export_operation(
+ body=body,
+ database_id=database_id,
+ project_id=project_id,
+ include_completed=False,
+ )
+ if existing_op is not None:
+ self.log.info(
+ "Reattaching to in-progress Firestore export operation
%s.",
+ existing_op.get("name"),
+ )
+ return existing_op
+
+ service = self.get_conn()
+ name = f"projects/{project_id}/databases/{database_id}"
+ try:
+ return (
+ service.projects()
+ .databases()
+ .exportDocuments(name=name, body=body)
+ .execute(num_retries=self.num_retries)
+ )
+ except HttpError as err:
+ if reattach_on_restart and err.resp.status in {400, 409}:
Review Comment:
**Major.** On a 400/409, a completed match is returned as success, but
nothing ties that operation to an earlier attempt of this task. The operations
list keeps recently completed exports, so this also matches an older run.
Concrete case: a scheduled Dag with a fixed `outputUriPrefix`, like the
how-to example in `example_firestore.py`.
1. Day 1 exports normally.
2. On day 2, the pre-check ignores completed operations and
`exportDocuments` returns 400 because the path exists.
3. This lookup finds day 1's operation, and the task succeeds with no new
export.
Today that run fails loudly. With this change, downstream tasks read stale
data and the task state shows nothing wrong.
The status code is also the only check on the error. An unrelated 400, such
as an invalid `snapshotTime`, is reported as success too when an old match
exists.
Suggestion: accept a completed operation only if it started after this Dag
run did. The operator can pass `context["dag_run"].start_date`, and the hook
can compare it with `metadata.startTime`. Alternatively, make completed-export
recovery opt-in and keep in-progress reattach as the default.
Either way, the `reattach_on_restart` docstrings on the hook and operator
need updating. They describe checking for completed operations before
submitting, but that check runs only after the HTTP error.
##########
providers/google/src/airflow/providers/google/firebase/hooks/firestore.py:
##########
@@ -84,9 +91,188 @@ def get_conn(self):
self._conn = build_from_document(non_authorized_conn._rootDesc,
http=http_authorized)
return self._conn
+ def get_operation(self, operation_name: str) -> dict[str, Any]:
+ """
+ Retrieve the current state of a long-running Firestore operation.
+
+ :param operation_name: The resource name of the operation.
+ """
+ service = self.get_conn()
+ return (
+ service.projects()
+ .databases()
+ .operations()
+ .get(name=operation_name)
+ .execute(num_retries=self.num_retries)
+ )
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def list_operations(
+ self, database_id: str = "(default)", project_id: str =
PROVIDE_PROJECT_ID
+ ) -> list[dict[str, Any]]:
+ """
+ List long-running operations on the specified Firestore database.
+
+ :param database_id: The Database ID.
+ :param project_id: Optional, Google Cloud Project project_id where the
database belongs.
+ If set to None or missing, the default project_id from the Google
Cloud connection is used.
+ """
+ service = self.get_conn()
+ name = f"projects/{project_id}/databases/{database_id}"
+ operations_resource = service.projects().databases().operations()
+ all_operations: list[dict[str, Any]] = []
+ page_token: str | None = None
+
+ while True:
+ list_kwargs: dict[str, Any] = {"name": name}
+ if page_token:
+ list_kwargs["pageToken"] = page_token
+ response =
operations_resource.list(**list_kwargs).execute(num_retries=self.num_retries)
+ if not isinstance(response, dict):
+ break
+ operations = response.get("operations", [])
+ if isinstance(operations, list):
+ all_operations.extend(operations)
+ next_page_token = response.get("nextPageToken")
+ if not isinstance(next_page_token, str) or not next_page_token:
+ break
+ page_token = next_page_token
+
+ return all_operations
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def find_matching_export_operation(
+ self,
+ body: dict[str, Any],
+ database_id: str = "(default)",
+ project_id: str = PROVIDE_PROJECT_ID,
+ include_completed: bool = True,
+ ) -> dict[str, Any] | None:
+ """
+ Find an existing running or completed export operation matching the
export request body.
+
+ :param body: The export request body.
+ :param database_id: The Database ID.
+ :param project_id: Optional, Google Cloud Project project_id where the
database belongs.
+ :param include_completed: Whether to match already completed
operations in addition to
+ in-progress operations.
+ """
+ target_uri = body.get("outputUriPrefix")
+ if not target_uri or not isinstance(target_uri, str):
+ return None
+ target_collections = sorted(body.get("collectionIds") or [])
+ target_namespaces = sorted(body.get("namespaceIds") or [])
+
+ completed_match: dict[str, Any] | None = None
+ for op in self.list_operations(database_id=database_id,
project_id=project_id):
+ if not isinstance(op, dict):
+ continue
+ metadata = op.get("metadata")
+ if not isinstance(metadata, dict):
+ metadata = {}
+
+ meta_type = metadata.get("@type", "")
+ if meta_type and not meta_type.endswith("ExportDocumentsMetadata"):
+ continue
+
+ op_state = metadata.get("operationState", "")
+ if op.get("error") or op_state in {"FAILED", "CANCELLED",
"CANCELLING"}:
+ continue
+
+ op_collections = sorted(metadata.get("collectionIds") or [])
+ if target_collections != op_collections:
+ continue
+
+ op_namespaces = sorted(metadata.get("namespaceIds") or [])
+ if target_namespaces != op_namespaces:
+ continue
+
+ response_dict = op.get("response")
+ op_uri = metadata.get("outputUriPrefix") or (
+ response_dict.get("outputUriPrefix") if
isinstance(response_dict, dict) else None
+ )
+ if not isinstance(op_uri, str) or op_uri.rstrip("/") !=
target_uri.rstrip("/"):
+ continue
+
+ is_done = bool(op.get("done"))
+ if not is_done:
+ return op
+ if include_completed and completed_match is None:
+ completed_match = op
+
+ return completed_match
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def start_export_documents(
+ self,
+ body: dict[str, Any],
+ database_id: str = "(default)",
+ project_id: str = PROVIDE_PROJECT_ID,
+ reattach_on_restart: bool = True,
+ ) -> dict[str, Any]:
+ """
+ Start a Firestore export or return an existing running/completed
operation when reattaching.
+
+ :param body: The request body.
+ :param database_id: The Database ID.
+ :param project_id: Optional, Google Cloud Project project_id where the
database belongs.
+ :param reattach_on_restart: If True, check for an existing in-progress
export operation
+ matching ``body`` before submitting a new export request, and
recover if the export
+ already completed on a prior attempt.
+ """
+ if reattach_on_restart:
+ existing_op = self.find_matching_export_operation(
Review Comment:
**Minor.** This pre-check runs on every attempt, including the first, and it
needs `datastore.operations.list`. The predefined import/export roles include
that permission. A custom role with only `datastore.databases.export` and
`datastore.operations.get` works today, but after upgrading it gets a 403 here
before any export is submitted.
Could a 403 on this lookup log a warning and carry on to the export? If you
prefer to keep it strict, this needs a changelog note. From
`providers/AGENTS.md`: "When a provider change needs a user-visible note
(typically a breaking change or important behavior change that warrants
explanation), update the provider's `docs/changelog.rst` directly in the same
PR".
##########
providers/google/tests/unit/google/firebase/hooks/test_firestore.py:
##########
@@ -131,6 +136,229 @@ def test_error_operation(self, _, get_conn_mock):
with pytest.raises(AirflowException, match="error"):
self.hook.export_documents(body=EXPORT_DOCUMENT_BODY,
project_id=TEST_PROJECT_ID)
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn")
+ def test_list_operations(self, get_conn_mock):
+ service_mock = get_conn_mock.return_value
+ mock_list =
service_mock.projects.return_value.databases.return_value.operations.return_value.list
+ second_operation = {"name": "operation-name-2"}
+ mock_list.return_value.execute.side_effect = [
+ {"operations": [TEST_OPERATION], "nextPageToken": "page-2"},
+ {"operations": [second_operation]},
+ ]
+
+ ops = self.hook.list_operations(project_id=TEST_PROJECT_ID)
+
+ assert ops == [TEST_OPERATION, second_operation]
+ assert mock_list.call_args_list == [
+
mock.call(name="projects/firestore--project-id/databases/(default)"),
+
mock.call(name="projects/firestore--project-id/databases/(default)",
pageToken="page-2"),
+ ]
+
+ @pytest.mark.parametrize(
+ ("body", "operations", "include_completed", "expected_name"),
+ [
+ pytest.param(
+ EXPORT_DOCUMENT_BODY,
+ [
+ "invalid-entry",
+ {
+ "name": "op-other-type",
+ "metadata": {"@type":
"type.googleapis.com/ImportDocumentsMetadata"},
+ },
+ {
+ "name": "op-failed",
+ "metadata": {
+ "@type":
"type.googleapis.com/ExportDocumentsMetadata",
+ "operationState": "FAILED",
+ "outputUriPrefix":
"gs://test-bucket/test-namespace",
+ "collectionIds": ["test-collection"],
+ },
+ },
+ {
+ "name": "op-wrong-collection",
+ "metadata": {
+ "outputUriPrefix":
"gs://test-bucket/test-namespace/",
+ "collectionIds": ["other"],
+ },
+ },
+ {
+ "name": "op-wrong-namespace",
+ "metadata": {
+ "outputUriPrefix":
"gs://test-bucket/test-namespace/",
+ "collectionIds": ["test-collection"],
+ "namespaceIds": ["ns1"],
+ },
+ },
+ {
+ "name": "op-wrong-uri",
+ "metadata": {
+ "outputUriPrefix": "gs://other-bucket/path",
+ "collectionIds": ["test-collection"],
+ },
+ },
+ {
+ "name": "op-completed",
+ "done": True,
+ "metadata": "non-dict-metadata",
+ "response": {"outputUriPrefix":
"gs://test-bucket/test-namespace"},
+ },
+ {
+ "name": "op-in-progress",
+ "done": False,
+ "metadata": {
+ "@type":
"type.googleapis.com/google.firestore.admin.v1.ExportDocumentsMetadata",
+ "operationState": "PROCESSING",
+ "outputUriPrefix":
"gs://test-bucket/test-namespace",
+ "collectionIds": ["test-collection"],
+ },
+ },
+ ],
+ True,
+ "op-in-progress",
+ id="prefers_in_progress_match",
+ ),
+ pytest.param(
+ EXPORT_DOCUMENT_BODY,
+ [
+ {
+ "name": "op-completed",
+ "done": True,
+ "metadata": {
+ "collectionIds": ["test-collection"],
+ },
+ "response": {"outputUriPrefix":
"gs://test-bucket/test-namespace"},
+ }
+ ],
+ True,
+ "op-completed",
+ id="matches_completed_operation_via_response_uri",
+ ),
+ pytest.param(
+ EXPORT_DOCUMENT_BODY,
+ [
+ {
+ "name": "op-completed",
+ "done": True,
+ "metadata": {
+ "collectionIds": ["test-collection"],
+ },
+ "response": {"outputUriPrefix":
"gs://test-bucket/test-namespace"},
+ }
+ ],
+ False,
+ None,
+ id="ignores_completed_when_include_completed_false",
+ ),
+ pytest.param(
+ {"collectionIds": ["test-collection"]},
+ [
+ {
+ "name": "op-done-no-target-uri",
+ "done": True,
+ "metadata": {"collectionIds": ["test-collection"]},
+ }
+ ],
+ True,
+ None,
+ id="returns_none_when_no_target_uri_specified",
+ ),
+ ],
+ )
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.list_operations")
+ def test_find_matching_export_operation(
+ self, mock_list_ops, body, operations, include_completed, expected_name
+ ):
+ mock_list_ops.return_value = operations
+ result = self.hook.find_matching_export_operation(
+ body=body, project_id=TEST_PROJECT_ID,
include_completed=include_completed
+ )
+ if expected_name is None:
+ assert result is None
+ else:
+ assert result is not None
+ assert result["name"] == expected_name
+
+ @mock.patch(
+
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.find_matching_export_operation"
+ )
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn")
+ def test_export_documents_reattaches_to_in_progress_operation(self,
get_conn_mock, mock_find_op):
+ service_mock = get_conn_mock.return_value
+ mock_export_documents =
service_mock.projects.return_value.databases.return_value.exportDocuments
+ mock_operation_get = (
+
service_mock.projects.return_value.databases.return_value.operations.return_value.get
+ )
+ mock_find_op.return_value = {"name": "existing-op", "done": False}
+ mock_operation_get.return_value.execute.return_value =
TEST_DONE_OPERATION
+
+ self.hook.export_documents(body=EXPORT_DOCUMENT_BODY,
project_id=TEST_PROJECT_ID)
+
+ mock_export_documents.assert_not_called()
+ mock_operation_get.assert_called_once_with(name="existing-op")
+
+ @mock.patch(
+
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.find_matching_export_operation"
+ )
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn")
+ def test_export_documents_skips_when_already_completed_on_http_400(self,
get_conn_mock, mock_find_op):
+ service_mock = get_conn_mock.return_value
+ mock_export_documents =
service_mock.projects.return_value.databases.return_value.exportDocuments
+ mock_operation_get = (
+
service_mock.projects.return_value.databases.return_value.operations.return_value.get
+ )
+ http_400 = HttpError(httplib2.Response({"status": 400}), b"Path
already exists")
+ mock_export_documents.return_value.execute.side_effect = http_400
+ mock_find_op.side_effect = [
+ None,
+ {
+ "name": "completed-op",
+ "done": True,
+ "response": {"outputUriPrefix": "x"},
+ },
+ ]
+
+ self.hook.export_documents(body=EXPORT_DOCUMENT_BODY,
project_id=TEST_PROJECT_ID)
+
+ mock_export_documents.assert_called_once()
+ mock_operation_get.assert_not_called()
+
+ @mock.patch(
+
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.find_matching_export_operation"
+ )
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn")
+ def
test_start_export_documents_recovers_on_http_400_race_in_progress(self,
get_conn_mock, mock_find_op):
+ service_mock = get_conn_mock.return_value
+ mock_export_documents =
service_mock.projects.return_value.databases.return_value.exportDocuments
+ http_400 = HttpError(httplib2.Response({"status": 400}), b"Operation
already in progress")
+ mock_export_documents.return_value.execute.side_effect = http_400
+ mock_find_op.side_effect = [None, {"name": "race-op", "done": False}]
+
+ op = self.hook.start_export_documents(body=EXPORT_DOCUMENT_BODY,
project_id=TEST_PROJECT_ID)
+
+ assert op == {"name": "race-op", "done": False}
+
+ @mock.patch(
+
"airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.find_matching_export_operation"
+ )
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreHook.get_conn")
+ def test_start_export_documents_reraises_http_error_when_no_match(self,
get_conn_mock, mock_find_op):
+ service_mock = get_conn_mock.return_value
+ mock_export_documents =
service_mock.projects.return_value.databases.return_value.exportDocuments
+ http_400 = HttpError(httplib2.Response({"status": 400}), b"Invalid
argument")
+ mock_export_documents.return_value.execute.side_effect = http_400
+ mock_find_op.return_value = None
+
+ with pytest.raises(HttpError):
+ self.hook.start_export_documents(
+ body=EXPORT_DOCUMENT_BODY, project_id=TEST_PROJECT_ID,
reattach_on_restart=False
Review Comment:
**Minor.** With `reattach_on_restart=False` the lookup is never reached, so
this test would still pass with the whole `except` block removed. The case in
the test name, reattach on with no match, isn't covered.
From the review instructions: "**Flag any changed or added behaviour without
a corresponding test, and flag tests that a reviewer cannot fail by reverting
the PR's change.**"
Dropping the kwarg and asserting `mock_find_op.call_count == 2` would cover
it. The `False` case could be a separate test asserting
`mock_find_op.assert_not_called()`.
##########
providers/google/tests/unit/google/firebase/hooks/test_firestore.py:
##########
@@ -256,3 +484,22 @@ def test_create_build(self, mock_get_conn,
mock_project_id):
== "The project id must be passed either as keyword project_id
parameter or as project_id extra in "
"Google Cloud connection definition. Both are not set!"
)
+
+
+class TestCloudFirestoreAsyncHook:
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.firebase.hooks.firestore.CloudFirestoreAsyncHook.get_sync_hook")
+ async def test_get_operation(self, mock_get_sync_hook):
+ sync_hook_mock = mock.MagicMock()
Review Comment:
**Nit.** From the review instructions: "**Flag any `mock.Mock()` or
`mock.MagicMock()` without `spec` or `autospec`.** Unspec'd mocks silently
accept any attribute access, hiding real bugs."
```suggestion
sync_hook_mock = mock.create_autospec(CloudFirestoreHook,
instance=True)
```
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
