Fix missing get_backup method for Dataproc Metastore (#20326)

Found during #19891 fixing (yay! MyPy actually found some real errors). The `get_backup` method was missing in Dataproc Metastore implementation.
apache · Dec 15, 2021 · 21b8661 · 21b8661
1 parent cdaa9a2
commit 21b8661
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 0 deletions.
diff --git a/airflow/providers/google/cloud/hooks/dataproc_metastore.py b/airflow/providers/google/cloud/hooks/dataproc_metastore.py
@@ -453,6 +453,53 @@ def get_service(
  )
  return result
 
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_backup(
+ self,
+ project_id: str,
+ region: str,
+ service_id: str,
+ backup_id: str,
+ retry: Optional[Retry] = None,
+ timeout: Optional[float] = None,
+ metadata: Sequence[Tuple[str, str]] = (),
+ ) -> Backup:
+ """
+ Get backup from a service.
+
+ :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+ :type project_id: str
+ :param region: Required. The ID of the Google Cloud region that the service belongs to.
+ :type region: str
+ :param service_id: Required. The ID of the metastore service, which is used as the final component of
+ the metastore service's name. This value must be between 2 and 63 characters long inclusive, begin
+ with a letter, end with a letter or number, and consist of alphanumeric ASCII characters or
+ hyphens.
+
+ This corresponds to the ``service_id`` field on the ``request`` instance; if ``request`` is
+ provided, this should not be set.
+ :type service_id: str
+ :param backup_id: Required. The ID of the metastore service backup to restore from
+ :type backup_id: str
+ :param retry: Designation of what errors, if any, should be retried.
+ :type retry: google.api_core.retry.Retry
+ :param timeout: The timeout for this request.
+ :type timeout: float
+ :param metadata: Strings which should be sent along with the request as metadata.
+ :type metadata: Sequence[Tuple[str, str]]
+ """
+ backup = f'projects/{project_id}/locations/{region}/services/{service_id}/backups/{backup_id}'
+ client = self.get_dataproc_metastore_client()
+ result = client.get_backup(
+ request={
+ 'name': backup,
+ },
+ retry=retry,
+ timeout=timeout,
+ metadata=metadata,
+ )
+ return result
+
  @GoogleBaseHook.fallback_to_default_project_id
  def list_backups(
  self,

diff --git a/tests/providers/google/cloud/hooks/test_dataproc_metastore.py b/tests/providers/google/cloud/hooks/test_dataproc_metastore.py
@@ -129,6 +129,24 @@ def test_create_service(self, mock_client) -> None:
  timeout=None,
  )
 
+ @mock.patch(DATAPROC_METASTORE_STRING.format("DataprocMetastoreHook.get_dataproc_metastore_client"))
+ def test_get_backup(self, mock_client) -> None:
+ self.hook.get_backup(
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ service_id=TEST_SERVICE_ID,
+ backup_id=TEST_BACKUP_ID,
+ )
+ mock_client.assert_called_once()
+ mock_client.return_value.get_backup.assert_called_once_with(
+ request=dict(
+ name=TEST_NAME_BACKUPS.format(TEST_PROJECT_ID, TEST_REGION, TEST_SERVICE_ID, TEST_BACKUP_ID),
+ ),
+ metadata=(),
+ retry=None,
+ timeout=None,
+ )
+
  @mock.patch(DATAPROC_METASTORE_STRING.format("DataprocMetastoreHook.get_dataproc_metastore_client"))
  def test_delete_backup(self, mock_client) -> None:
  self.hook.delete_backup(