Repository: astronomer/astro-sdk

Issue:
Support for Google Drive as File source
Create support for google drive as a file location. A file(s) kept in google drive would be now loaded into the databases and can perform operations on it. 
For this, we use the Google API client to connect and retrieve the file.

One endpoint is used:

* [`files().list()`](https://developers.google.com/drive/api/v3/reference/files/list) to find files.

To use those file APIs you need to kind of visualise traversing the folders in a GUI. To get a file by path, you need to see what items are available in a folder and select the one with the name you want. If the item is a folder, you need to run the list against it again to see what items are in it. Not very efficient, but there are not many choices.



Discussion:
A few thoughts from initial observations and investigations:

1. Airflow’s GoogleDriveHook is pretty lacking at the moment[^1]. So we either need to contribute to that first, or just roll our own implementation in the SDK.
2. There’s a new Python package called [fsspec](https://filesystem-spec.readthedocs.io/) that provides file API for Google Drive (and some other storages). It is gaining populatiry recently.
3. We are currently using `smart_open`, which lacks Google Drive support. So we need to implement a [custom transport](https://github.com/RaRe-Technologies/smart_open/blob/master/extending.md) for it.
4. But fsspec actully can also do `open`, so it’s entirely possible to replace `smart_open` with fsspec for our purposes. And since fsspec has built-in Google Drive support, we’ll be able to add it more easily.

I’m wondering how it’s best to approach this. We can either replace `smart_open` with fsspec, or implement a custom `smart_open` transport. If we decide to go with the latter, we can also try to implement the custom transport directly with Google API client, or add fsspec as a dependency.

[^1]: Currently only methods to check, read, and write a file.
I took a look at fsspec’s Google Drive implementation and it seems easy enough to just use the API client, so I’m going to do that for now.

I wonder if there’s a good way to test the classes. A Moto equivalent for googleapiclient would be useful…
PR to add the file location has been created (see ref link above)

Resolving patch:
diff --git a/python-sdk/example_dags/example_load_file.py b/python-sdk/example_dags/example_load_file.py
--- a/python-sdk/example_dags/example_load_file.py
+++ b/python-sdk/example_dags/example_load_file.py
@@ -1,3 +1,21 @@
+"""
+Pre-requisites for load_file_example_19:
+ - Install dependencies for Astro Python SDK with Google, refer to README.md
+ - You can either specify a service account key file and set `GOOGLE_APPLICATION_CREDENTIALS`
+    with the file path to the service account.
+ - In the connection we need to specfiy the scopes.
+    Connection variable is ``extra__google_cloud_platform__scope``
+    or in Airflow Connections UI ``Scopes (comma separated)``
+    For ex:- https://www.googleapis.com/auth/drive.readonly
+    Please refer to https://developers.google.com/identity/protocols/oauth2/scopes#drive for more details.
+ - In Google Cloud, for the project we need to enable the Google Drive API.
+    To enable the API please refer https://developers.google.com/drive/api/guides/enable-drive-api
+ - Create a Google Drive folder (or) use an existing folder with a file inside it,
+    and share the file with service account email id in order for it to be able to access those
+    folders/files. In this example DAG, we will load this file into Snowflake table.
+    For sharing a file/folder
+    please refer https://www.labnol.org/google-api-service-account-220404#4-share-a-drive-folder
+"""
 import os
 import pathlib
 from datetime import datetime, timedelta
@@ -13,6 +31,7 @@
 # To create IAM role with needed permissions,
 # refer: https://www.dataliftoff.com/iam-roles-for-loading-data-from-s3-into-redshift/
 REDSHIFT_NATIVE_LOAD_IAM_ROLE_ARN = os.getenv("REDSHIFT_NATIVE_LOAD_IAM_ROLE_ARN")
+SNOWFLAKE_CONN_ID = "snowflake_conn"
 
 CWD = pathlib.Path(__file__).parent
 default_args = {
@@ -221,4 +240,17 @@
     )
     # [END load_file_example_18]
 
+    # [START load_file_example_19]
+    aql.load_file(
+        input_file=File(path="gdrive://test-google-drive-support/sample.csv", conn_id="gdrive_conn"),
+        output_table=Table(
+            conn_id=SNOWFLAKE_CONN_ID,
+            metadata=Metadata(
+                database=os.environ["SNOWFLAKE_DATABASE"],
+                schema=os.environ["SNOWFLAKE_SCHEMA"],
+            ),
+        ),
+    )
+    # [END load_file_example_19]
+
     aql.cleanup()
diff --git a/python-sdk/src/astro/constants.py b/python-sdk/src/astro/constants.py
--- a/python-sdk/src/astro/constants.py
+++ b/python-sdk/src/astro/constants.py
@@ -20,6 +20,7 @@ class FileLocation(Enum):
     HTTP = "http"
     HTTPS = "https"
     GS = "gs"  # Google Cloud Storage
+    GOOGLE_DRIVE = "gdrive"
     S3 = "s3"  # Amazon S3
     # [END filelocation]
 
diff --git a/python-sdk/src/astro/files/locations/google/gdrive.py b/python-sdk/src/astro/files/locations/google/gdrive.py
new file mode 100644
--- /dev/null
+++ b/python-sdk/src/astro/files/locations/google/gdrive.py
@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import pathlib
+import urllib.parse
+from typing import Any, Iterator, Sequence
+
+from airflow.providers.google.suite.hooks.drive import GoogleDriveHook
+
+from astro.constants import FileLocation
+from astro.files.locations.base import BaseFileLocation
+
+
+def _quote(s: str) -> str:
+    s = s.replace("'", "\\'")  # Escape single quote characters.
+    return f"'{s}'"  # Quote.
+
+
+def _find_item_id(path_parts: str, *, conn: Any) -> str:
+    """Drill down the directory hierarchy to find the ID of an item.
+
+    Google Drive does not have a very good random access API by path, so we
+    must look at each directory, level by level, to find the item we want.
+    """
+    current_parent = "root"
+    folders = path_parts.split("/")
+    # First tries to enter directories
+    for current_folder in folders:
+        conditions = [
+            "mimeType = 'application/vnd.google-apps.folder'",
+            f"name='{current_folder}'",
+        ]
+        if current_parent != "root":
+            conditions.append(f"'{current_parent}' in parents")
+        result = (
+            conn.files().list(q=" and ".join(conditions), spaces="drive", fields="files(id, name)").execute()
+        )
+        files = result.get("files", [])
+        if not files:
+            # If the directory does not exist, break loops
+            break
+        current_parent = files[0].get("id")
+    return current_parent
+
+
+def _find_contains(folder_id: str, pattern: str, *, conn: Any) -> Iterator[str]:
+    """Find files in a directory matching given pattern.
+
+    This uses *contains* to search. See Google Drive documentation to
+    understand the implications:
+    https://developers.google.com/drive/api/guides/ref-search-terms
+    """
+    page_token = None
+    while True:
+        response = (
+            conn.files()
+            .list(
+                q=f"name contains {_quote(pattern)} and {_quote(folder_id)} in parents ",
+                fields="nextPageToken, files(name,id, mimeType,webContentLink,size)",
+                pageToken=page_token,
+            )
+            .execute()
+        )
+        yield from (f["webContentLink"] for f in response["files"])
+        page_token = response.get("nextPageToken", None)
+        if page_token is None:
+            return
+
+
+class GdriveLocation(BaseFileLocation):
+    """Handler for Google Drive operators."""
+
+    location_type = FileLocation.GOOGLE_DRIVE
+
+    @property
+    def hook(self) -> GoogleDriveHook:
+        if self.conn_id:
+            return GoogleDriveHook(gcp_conn_id=self.conn_id)
+        return GoogleDriveHook()
+
+    @property
+    def transport_params(self) -> dict:
+        """Get Google Drive client."""
+        return {"client": self.hook.get_conn()}
+
+    @property
+    def openlineage_dataset_namespace(self) -> str:
+        """
+        Returns the open lineage dataset namespace as per
+        https://github.com/OpenLineage/OpenLineage/blob/main/spec/Naming.md
+        """
+        parsed_url = urllib.parse.urlsplit(self.path)
+        return f"{parsed_url.scheme}://{parsed_url.netloc}"
+
+    @property
+    def openlineage_dataset_name(self) -> str:
+        """
+        Returns the open lineage dataset name as per
+        https://github.com/OpenLineage/OpenLineage/blob/main/spec/Naming.md
+        """
+        return urllib.parse.urlsplit(self.path).path
+
+    def _get_rel_path_parts(self) -> Sequence[str]:
+        p = urllib.parse.urlsplit(self.path).path.lstrip("/")
+        return pathlib.PurePosixPath(p).parts
+
+    def _get_conn(self) -> Any:
+        return self.hook.get_conn()
+
+    @property
+    def paths(self) -> list[str]:
+        path_parts = self._get_rel_path_parts()
+        if not path_parts:
+            return []
+
+        conn = self._get_conn()
+        url = urllib.parse.urlsplit(self.path)
+        result_paths = []
+        try:
+            folder_id = _find_item_id(url.netloc + url.path, conn=conn)
+        except FileNotFoundError:
+            raise FileNotFoundError(self.path) from None
+        for name in _find_contains(folder_id, path_parts[-1], conn=conn):
+            result_paths.append(name)
+        return result_paths
+
+    @property
+    def size(self) -> int:
+        path_parts = self._get_rel_path_parts()
+        conn = self._get_conn()
+        url = urllib.parse.urlsplit(self.path)
+        try:
+            folder_id = _find_item_id(url.netloc + url.path, conn=conn)
+        except FileNotFoundError:
+            raise FileNotFoundError(self.path) from None
+
+        response = (
+            conn.files()
+            .list(
+                q=f"name contains {_quote(path_parts[-1])} and {_quote(folder_id)} in parents ",
+                fields="files(name,id,size)",
+            )
+            .execute()
+        )
+        value: int = 0
+        try:
+            for f in response["files"]:
+                value = int(f["size"])
+        except KeyError:
+            raise IsADirectoryError(self.path) from None
+        return value


Test patch:
diff --git a/python-sdk/tests/files/locations/test_google_drive.py b/python-sdk/tests/files/locations/test_google_drive.py
new file mode 100644
--- /dev/null
+++ b/python-sdk/tests/files/locations/test_google_drive.py
@@ -0,0 +1,95 @@
+from unittest.mock import patch
+
+import pytest
+from airflow.providers.google.suite.hooks.drive import GoogleDriveHook
+from googleapiclient.discovery import Resource
+
+from astro.files.locations import create_file_location
+
+
+def test_get_transport_params_for_gdrive():  # skipcq: PYL-W0612, PTC-W0065
+    """test get_transport_params() method which should return gdrive resource client"""
+    path = "gdrive://bucket/some-file"
+    location = create_file_location(path)
+    credentials = location.transport_params
+    assert isinstance(credentials["client"], Resource)
+
+
+@patch("airflow.providers.google.suite.hooks.drive.GoogleDriveHook.get_conn")
+@patch("astro.files.locations.google.gdrive._find_item_id")
+def test_get_paths_from_gdrive(
+    mock_folder_id,
+    mock_get_conn,
+):
+    """Get the list of files from the gdrive path"""
+    mock_get_conn.return_value.files.return_value.list.return_value.execute.return_value = {
+        "files": [{"webContentLink": "ADOPTION_CENTER_1_unquoted.csv"}]
+    }
+    mock_folder_id.return_value = "root"
+    location = create_file_location("gdrive://data/ADOPTION_CENTER_1_unquoted.csv")
+    assert sorted(location.paths) == sorted(["ADOPTION_CENTER_1_unquoted.csv"])
+
+
+@patch("airflow.providers.google.suite.hooks.drive.GoogleDriveHook.get_conn")
+@patch("astro.files.locations.google.gdrive._find_item_id")
+def test_size(mock_folder_id, mock_get_conn):
+    """Test get_size() of for Google Drive file."""
+    mock_get_conn.return_value.files.return_value.list.return_value.execute.return_value = {
+        "files": [{"size": "110"}]
+    }
+    mock_folder_id.return_value = "root"
+    path = "gdrive://data/ADOPTION_CENTER_1_unquoted.csv"
+    location = create_file_location(path)
+    assert location.size == 110
+
+
+def test_hook():
+    """Test whether GoogleDriveHook is being called or not."""
+    path = "gdrive://bucket/some-file"
+    location = create_file_location(path)
+    hook = location.hook
+    assert isinstance(hook, GoogleDriveHook)
+
+
+@patch("astro.files.locations.google.gdrive.GdriveLocation._get_rel_path_parts")
+def test_gdrive_file_not_found(
+    mock_get_path,
+):
+    """Test when no file is found in gdrive path"""
+    mock_get_path.return_value = []
+    location = create_file_location("gdrive://bucket/some-file")
+    assert location.paths == []
+
+
+@patch("astro.files.locations.google.gdrive._find_item_id")
+def test_remote_object_exception(
+    mock_folder_id,
+):
+    """Raise exception when gdrive filepath doesn't exists"""
+    mock_folder_id.side_effect = FileNotFoundError()
+    location = create_file_location("gdrive://data/ADOPTION_CENTER_1_unquoted.csv")
+    with pytest.raises(FileNotFoundError):
+        location.paths
+
+
+@patch("astro.files.locations.google.gdrive._find_item_id")
+def test_size_exception(mock_folder_id):
+    """Raise exception for a file not existing in gdrive filepath"""
+    mock_folder_id.side_effect = FileNotFoundError()
+    location = create_file_location("gdrive://data/ADOPTION_CENTER_1_unquoted.csv")
+    with pytest.raises(FileNotFoundError):
+        location.size
+
+
+@patch("airflow.providers.google.suite.hooks.drive.GoogleDriveHook.get_conn")
+@patch("astro.files.locations.google.gdrive._find_item_id")
+def test_size_key_error(mock_folder_id, mock_get_conn):
+    """Raise exception when gdrive list doesn't return size attribute"""
+    mock_get_conn.return_value.files.return_value.list.return_value.execute.return_value = {
+        "files": [{"sizes": "110"}]
+    }
+    mock_folder_id.return_value = "root"
+    path = "gdrive://data/ADOPTION_CENTER_1_unquoted.csv"
+    location = create_file_location(path)
+    with pytest.raises(IsADirectoryError):
+        location.size
diff --git a/python-sdk/tests/files/locations/test_location_base.py b/python-sdk/tests/files/locations/test_location_base.py
--- a/python-sdk/tests/files/locations/test_location_base.py
+++ b/python-sdk/tests/files/locations/test_location_base.py
@@ -55,6 +55,7 @@ class TestLocation:  # skipcq: PY-D0002
         (FileLocation.HTTPS, "https://domain/some-file", "https://domain"),
         (FileLocation.S3, "s3://bucket/some-file", "s3://bucket"),
         (FileLocation.GS, "gs://bucket/some-file", "gs://bucket"),
+        (FileLocation.GOOGLE_DRIVE, "gdrive://bucket/some-file", "gdrive://bucket"),
     ],
 )
 def test_openlineage_file_dataset_namespace(file_location, filepath, namespace):
@@ -74,6 +75,7 @@ def test_openlineage_file_dataset_namespace(file_location, filepath, namespace):
         (FileLocation.HTTPS, "https://domain/some-file", "/some-file"),
         (FileLocation.S3, "s3://bucket/some-file", "/some-file"),
         (FileLocation.GS, "gs://bucket/some-file", "/some-file"),
+        (FileLocation.GOOGLE_DRIVE, "gdrive://bucket/some-file", "/some-file"),
     ],
 )
 def test_openlineage_file_dataset_name(file_location, filepath, dataset_name):
diff --git a/python-sdk/tests/test_constants.py b/python-sdk/tests/test_constants.py
--- a/python-sdk/tests/test_constants.py
+++ b/python-sdk/tests/test_constants.py
@@ -2,7 +2,7 @@
 
 
 def test_supported_file_locations():
-    expected = ["gs", "http", "https", "local", "s3"]
+    expected = ["gdrive", "gs", "http", "https", "local", "s3"]
     assert sorted(SUPPORTED_FILE_LOCATIONS) == expected