Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 19 additions & 3 deletions airflow/providers/google/suite/hooks/drive.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,31 +128,47 @@ def get_media_request(self, file_id: str) -> HttpRequest:
request = service.files().get_media(fileId=file_id)
return request

def exists(self, folder_id: str, file_name: str, drive_id: Optional[str] = None):
def exists(
self, folder_id: str, file_name: str, drive_id: Optional[str] = None, *, include_trashed: bool = True
):
"""
Checks to see if a file exists within a Google Drive folder

:param folder_id: The id of the Google Drive folder in which the file resides
:param file_name: The name of a file in Google Drive
:param drive_id: Optional. The id of the shared Google Drive in which the file resides.
:param include_trashed: Whether to include objects in trash or not, default True as in Google API.

:return: True if the file exists, False otherwise
:rtype: bool
"""
return bool(self.get_file_id(folder_id=folder_id, file_name=file_name, drive_id=drive_id))
return bool(
self.get_file_id(
folder_id=folder_id, file_name=file_name, include_trashed=include_trashed, drive_id=drive_id
)
)

def get_file_id(self, folder_id: str, file_name: str, drive_id: Optional[str] = None):
def get_file_id(
self, folder_id: str, file_name: str, drive_id: Optional[str] = None, *, include_trashed: bool = True
):
"""
Returns the file id of a Google Drive file

:param folder_id: The id of the Google Drive folder in which the file resides
:param file_name: The name of a file in Google Drive
:param drive_id: Optional. The id of the shared Google Drive in which the file resides.
:param include_trashed: Whether to include objects in trash or not, default True as in Google API.

:return: Google Drive file id if the file exists, otherwise None
:rtype: str if file exists else None
"""
query = f"name = '{file_name}'"
if folder_id:
query += f" and parents in '{folder_id}'"

if not include_trashed:
query += " and trashed=false"

service = self.get_conn()
if drive_id:
files = (
Expand Down
23 changes: 21 additions & 2 deletions tests/providers/google/suite/hooks/test_drive.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,9 @@ def test_exists_when_file_exists(self, mock_get_conn, mock_method):
file_name = "abc123.csv"

result_value = self.gdrive_hook.exists(folder_id=folder_id, file_name=file_name, drive_id=drive_id)
mock_method.assert_called_once_with(folder_id=folder_id, file_name=file_name, drive_id=drive_id)
mock_method.assert_called_once_with(
folder_id=folder_id, file_name=file_name, drive_id=drive_id, include_trashed=True
)
self.assertEqual(True, result_value)

@mock.patch("airflow.providers.google.suite.hooks.drive.GoogleDriveHook.get_file_id")
Expand All @@ -178,7 +180,24 @@ def test_exists_when_file_not_exists(self, mock_get_conn, mock_method):
file_name = "abc123.csv"

self.gdrive_hook.exists(folder_id=folder_id, file_name=file_name, drive_id=drive_id)
mock_method.assert_called_once_with(folder_id=folder_id, file_name=file_name, drive_id=drive_id)
mock_method.assert_called_once_with(
folder_id=folder_id, file_name=file_name, drive_id=drive_id, include_trashed=True
)

@mock.patch("airflow.providers.google.suite.hooks.drive.GoogleDriveHook.get_file_id")
@mock.patch("airflow.providers.google.suite.hooks.drive.GoogleDriveHook.get_conn")
def test_exists_when_trashed_is_false(self, mock_get_conn, mock_method):
folder_id = "abxy1z"
drive_id = "abc123"
file_name = "abc123.csv"
include_trashed = False

self.gdrive_hook.exists(
folder_id=folder_id, file_name=file_name, drive_id=drive_id, include_trashed=include_trashed
)
mock_method.assert_called_once_with(
folder_id=folder_id, file_name=file_name, drive_id=drive_id, include_trashed=False
)

@mock.patch("airflow.providers.google.suite.hooks.drive.GoogleDriveHook.get_conn")
def test_get_media_request(self, mock_get_conn):
Expand Down