airbyte-source-google-drive 0.1.1__tar.gz → 0.2.0__tar.gz

{airbyte_source_google_drive-0.1.1 → airbyte_source_google_drive-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: airbyte-source-google-drive
-Version: 0.1.1
+Version: 0.2.0
 Summary: Source implementation for Google Drive.
 License: ELv2
 Author: Airbyte
@@ -10,7 +10,7 @@ Classifier: License :: Other/Proprietary License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: airbyte-cdk[file-based] (>=6.18.2,<7.0.0)
+Requires-Dist: airbyte-cdk[file-based] (>=6.33.6,<7.0.0)
 Requires-Dist: google-api-python-client (==2.104.0)
 Requires-Dist: google-api-python-client-stubs (==1.18.0)
 Requires-Dist: google-auth-httplib2 (==0.1.1)

{airbyte_source_google_drive-0.1.1 → airbyte_source_google_drive-0.2.0}/pyproject.toml

@@ -5,7 +5,7 @@ requires = [
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
-version = "0.1.1"
+version = "0.2.0"
 name = "airbyte-source-google-drive"
 description = "Source implementation for Google Drive."
 authors = [
@@ -31,7 +31,7 @@ google-api-python-client-stubs = "==1.18.0"
 extras = [
     "file-based",
 ]
-version = "^6.18.2"
+version = "^6.33.6"
 
 [tool.poetry.scripts]
 source-google-drive = "source_google_drive.run:run"
@@ -39,3 +39,8 @@ source-google-drive = "source_google_drive.run:run"
 [tool.poetry.group.dev.dependencies]
 pytest-mock = "^3.12.0"
 pytest = "^7.4"
+
+[tool.poe]
+include = [
+    "${POE_GIT_DIR}/poe-tasks/poetry-connector-tasks.toml",
+]

airbyte_source_google_drive-0.2.0/source_google_drive/schemas/file_permissions.json

@@ -0,0 +1,12 @@
+{
+  "type": "object",
+  "properties": {
+    "id": { "type": "string" },
+    "file_path": { "type": "string" },
+    "allowed_identity_remote_ids": {
+      "type": "array",
+      "items": { "type": "string" }
+    },
+    "publicly_accessible": { "type": "boolean" }
+  }
+}

airbyte_source_google_drive-0.2.0/source_google_drive/schemas/identities.json

@@ -0,0 +1,14 @@
+{
+  "type": "object",
+  "properties": {
+    "id": { "type": "string" },
+    "remote_id": { "type": "string" },
+    "parent_id": { "type": ["null", "string"] },
+    "name": { "type": ["null", "string"] },
+    "description": { "type": ["null", "string"] },
+    "email_address": { "type": ["null", "string"] },
+    "member_email_addresses": { "type": ["null", "array"] },
+    "type": { "type": "string" },
+    "modified_at": { "type": "string" }
+  }
+}

{airbyte_source_google_drive-0.1.1 → airbyte_source_google_drive-0.2.0}/source_google_drive/source.py

@@ -6,7 +6,7 @@
 from typing import Any, Mapping, Optional
 
 from airbyte_cdk import AdvancedAuth, ConfiguredAirbyteCatalog, ConnectorSpecification, OAuthConfigSpecification, TState
-from airbyte_cdk.models import AuthFlowType
+from airbyte_cdk.models import AuthFlowType, OauthConnectorInputSpecification
 from airbyte_cdk.sources.file_based.file_based_source import FileBasedSource
 from airbyte_cdk.sources.file_based.stream.cursor.default_file_based_cursor import DefaultFileBasedCursor
 from source_google_drive.spec import SourceGoogleDriveSpec
@@ -28,6 +28,11 @@ class SourceGoogleDrive(FileBasedSource):
         """
         Returns the specification describing what fields can be configured by a user when setting up a file-based source.
         """
+        oauth_connector_input_specification = OauthConnectorInputSpecification(
+            consent_url="https://accounts.google.com/o/oauth2/v2/auth?{{client_id_param}}&{{redirect_uri_param}}&response_type=code&{{scope_param}}&access_type=offline&{{state_param}}&include_granted_scopes=true&prompt=consent",
+            access_token_url="https://oauth2.googleapis.com/token?{{client_id_param}}&{{client_secret_param}}&{{auth_code_param}}&{{redirect_uri_param}}&grant_type=authorization_code",
+            scope="https://www.googleapis.com/auth/drive.readonly https://www.googleapis.com/auth/admin.directory.group.readonly https://www.googleapis.com/auth/admin.directory.group.member.readonly https://www.googleapis.com/auth/admin.directory.user.readonly",
+        )
 
         return ConnectorSpecification(
             documentationUrl=self.spec_class.documentation_url(),
@@ -37,10 +42,17 @@ class SourceGoogleDrive(FileBasedSource):
                 predicate_key=["credentials", "auth_type"],
                 predicate_value="Client",
                 oauth_config_specification=OAuthConfigSpecification(
+                    oauth_connector_input_specification=oauth_connector_input_specification,
                     complete_oauth_output_specification={
                         "type": "object",
                         "additionalProperties": False,
-                        "properties": {"refresh_token": {"type": "string", "path_in_connector_config": ["credentials", "refresh_token"]}},
+                        "properties": {
+                            "refresh_token": {
+                                "type": "string",
+                                "path_in_connector_config": ["credentials", "refresh_token"],
+                                "path_in_oauth_response": ["refresh_token"],
+                            }
+                        },
                     },
                     complete_oauth_server_input_specification={
                         "type": "object",

{airbyte_source_google_drive-0.1.1 → airbyte_source_google_drive-0.2.0}/source_google_drive/spec.py

@@ -1,15 +1,52 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
-
-from typing import Any, Dict, Literal, Union
+import uuid
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, Literal, Optional, Union
 
 import dpath.util
 from pydantic.v1 import BaseModel, Field
 
 from airbyte_cdk import OneOfOptionConfig
-from airbyte_cdk.sources.file_based.config.abstract_file_based_spec import AbstractFileBasedSpec, DeliverRawFiles, DeliverRecords
+from airbyte_cdk.sources.file_based.config.abstract_file_based_spec import (
+    AbstractFileBasedSpec,
+    DeliverRawFiles,
+    DeliverRecords,
+)
+from airbyte_cdk.sources.file_based.config.abstract_file_based_spec import (
+    DeliverPermissions as DeliverPermissionsBase,
+)
+
+
+class RemoteIdentityType(Enum):
+    USER = "user"
+    GROUP = "group"
+
+
+class RemoteIdentity(BaseModel):
+    id: uuid.UUID
+    remote_id: str
+    parent_id: str | None = None
+    name: str | None = None
+    description: str | None = None
+    email_address: str | None = None
+    member_email_addresses: list[str] | None = None
+    type: RemoteIdentityType
+    modified_at: datetime
+
+
+class RemotePermissions(BaseModel):
+    id: str
+    file_path: str
+    allowed_identity_remote_ids: list[str] | None = None
+    denied_identity_remote_ids: list[str] | None = None
+    publicly_accessible: bool = False
+
+
+class DeliverPermissions(DeliverPermissionsBase):
+    domain: Optional[str] = Field(title="Domain", description="The Google domain of the identities.", airbyte_hidden=False, order=1)
 
 
 class OAuthCredentials(BaseModel):
@@ -60,7 +97,7 @@ class SourceGoogleDriveSpec(AbstractFileBasedSpec, BaseModel):
         pattern_descriptor="https://drive.google.com/drive/folders/MY-FOLDER-ID",
     )
 
-    delivery_method: DeliverRecords | DeliverRawFiles = Field(
+    delivery_method: DeliverRecords | DeliverRawFiles | DeliverPermissions = Field(
         title="Delivery Method",
         discriminator="delivery_type",
         type="object",

{airbyte_source_google_drive-0.1.1 → airbyte_source_google_drive-0.2.0}/source_google_drive/stream_reader.py

@@ -6,11 +6,13 @@
 import io
 import json
 import logging
+import uuid
 from datetime import datetime
 from io import IOBase
 from os.path import getsize
-from typing import Dict, Iterable, List, Optional, Set
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Set, Tuple
 
+import pytz
 from google.oauth2 import credentials, service_account
 from googleapiclient.discovery import build
 from googleapiclient.http import MediaIoBaseDownload
@@ -19,10 +21,12 @@ from airbyte_cdk import AirbyteTracedException, FailureType
 from airbyte_cdk.sources.file_based.exceptions import FileSizeLimitError
 from airbyte_cdk.sources.file_based.file_based_stream_reader import AbstractFileBasedStreamReader, FileReadMode
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
+from airbyte_cdk.sources.streams.core import package_name_from_class
+from airbyte_cdk.sources.utils.schema_helpers import InternalConfig, ResourceSchemaLoader
 from source_google_drive.utils import get_folder_id
 
 from .exceptions import ErrorDownloadingFile, ErrorFetchingMetadata
-from .spec import SourceGoogleDriveSpec
+from .spec import RemoteIdentity, RemoteIdentityType, RemotePermissions, SourceGoogleDriveSpec
 
 
 FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
@@ -54,6 +58,24 @@ DOWNLOADABLE_DOCUMENTS_MIME_TYPES = {
     GOOGLE_DRAWING_MIME_TYPE: {EXPORT_MEDIA_MIME_TYPE_KEY: EXPORT_MEDIA_MIME_TYPE_PDF, DOCUMENT_FILE_EXTENSION_KEY: ".pdf"},
 }
 
+PUBLIC_PERMISSION_IDS = [
+    "anyoneWithLink",
+    "anyoneCanFind",
+    "domainCanFind",
+    "domainWithLink",
+]
+
+
+DRIVE_SERVICE_SCOPES = [
+    "https://www.googleapis.com/auth/admin.directory.group.readonly",
+    "https://www.googleapis.com/auth/admin.directory.group.member.readonly",
+    "https://www.googleapis.com/auth/admin.directory.user.readonly",
+]
+
+
+def datetime_now() -> datetime:
+    return datetime.now(pytz.UTC)
+
 
 class GoogleDriveRemoteFile(RemoteFile):
     id: str
@@ -68,6 +90,7 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
     def __init__(self):
         super().__init__()
         self._drive_service = None
+        self._directory_service = None
 
     @property
     def config(self) -> SourceGoogleDriveSpec:
@@ -87,29 +110,41 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
         assert isinstance(value, SourceGoogleDriveSpec)
         self._config = value
 
-    @property
-    def google_drive_service(self):
+    def _build_google_service(self, service_name: str, version: str, scopes: List[str] = None):
         if self.config is None:
             # We shouldn't hit this; config should always get set before attempting to
             # list or read files.
-            raise ValueError("Source config is missing; cannot create the Google Drive client.")
+            raise ValueError(f"Source config is missing; cannot create the Google {service_name} client.")
         try:
-            if self._drive_service is None:
-                if self.config.credentials.auth_type == "Client":
-                    creds = credentials.Credentials.from_authorized_user_info(self.config.credentials.dict())
-                else:
-                    creds = service_account.Credentials.from_service_account_info(json.loads(self.config.credentials.service_account_info))
-                self._drive_service = build("drive", "v3", credentials=creds)
+            if self.config.credentials.auth_type == "Client":
+                creds = credentials.Credentials.from_authorized_user_info(self.config.credentials.dict())
+            else:
+                creds = service_account.Credentials.from_service_account_info(
+                    json.loads(self.config.credentials.service_account_info), scopes=scopes
+                )
+            google_service = build(service_name, version, credentials=creds)
         except Exception as e:
             raise AirbyteTracedException(
                 internal_message=str(e),
-                message="Could not authenticate with Google Drive. Please check your credentials.",
+                message=f"Could not authenticate with Google {service_name}. Please check your credentials.",
                 failure_type=FailureType.config_error,
                 exception=e,
             )
 
+        return google_service
+
+    @property
+    def google_drive_service(self):
+        if self._drive_service is None:
+            self._drive_service = self._build_google_service("drive", "v3")
         return self._drive_service
 
+    @property
+    def google_directory_service(self):
+        if self._directory_service is None:
+            self._directory_service = self._build_google_service("admin", "directory_v1", DRIVE_SERVICE_SCOPES)
+        return self._directory_service
+
     def get_matching_files(self, globs: List[str], prefix: Optional[str], logger: logging.Logger) -> Iterable[RemoteFile]:
         """
         Get all files matching the specified glob patterns.
@@ -282,3 +317,137 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
 
         except Exception as e:
             raise ErrorDownloadingFile(f"There was an error while trying to download the file {file.uri}: {str(e)}")
+
+    def get_file_permissions(self, file_id: str, file_name: str, logger: logging.Logger) -> Tuple[List[RemoteIdentity], bool]:
+        """
+        Retrieves the permissions of a file in Google Drive and checks for public access.
+
+        Args:
+            file_id (str): The file to get permissions for.
+            file_name (str): The name of the file to get permissions for.
+            logger (logging.Logger): Logger for debugging and information.
+
+        Returns:
+            Tuple(List[RemoteFileIdentity], boolean): A list of RemoteFileIdentity objects containing permission details.
+        """
+        try:
+            request = self.google_drive_service.permissions().list(
+                fileId=file_id,
+                fields="permissions, permissions/role, permissions/type, permissions/id, permissions/emailAddress",
+                supportsAllDrives=True,
+            )
+            response = request.execute()
+            permissions = response.get("permissions", [])
+            is_public = False
+
+            remote_identities = []
+
+            for p in permissions:
+                identity = self._to_remote_file_identity(p)
+                if p.get("id") in PUBLIC_PERMISSION_IDS:
+                    is_public = True
+                if identity is not None:
+                    remote_identities.append(identity)
+
+            return remote_identities, is_public
+        except Exception as e:
+            raise ErrorFetchingMetadata(f"An error occurred while retrieving file permissions: {str(e)}")
+
+    def _to_remote_file_identity(self, identity: dict[str, Any]) -> RemoteIdentity | None:
+        if identity.get("id") in PUBLIC_PERMISSION_IDS:
+            return None
+        if identity.get("deleted") is True:
+            return None
+
+        return RemoteIdentity(
+            modified_at=datetime.now(),
+            id=uuid.uuid4(),
+            remote_id=identity.get("emailAddress"),
+            name=identity.get("name"),
+            email_address=identity.get("emailAddress"),
+            type=identity.get("type"),
+            description=None,
+        )
+
+    def get_file_acl_permissions(self, file: GoogleDriveRemoteFile, logger: logging.Logger) -> Dict[str, Any]:
+        remote_identities, is_public = self.get_file_permissions(file.id, file_name=file.uri, logger=logger)
+        return RemotePermissions(
+            id=file.id,
+            file_path=file.uri,
+            allowed_identity_remote_ids=[p.remote_id for p in remote_identities],
+            publicly_accessible=is_public,
+        ).dict(exclude_none=True)
+
+    def _get_looping_google_api_list_response(
+        self, service: Any, key: str, args: dict[str, Any], logger: logging.Logger
+    ) -> Iterator[dict[str, Any]]:
+        try:
+            looping = True
+            next_page_token: str | None = None
+            while looping:
+                rsp = service.list(pageToken=next_page_token, **args).execute()
+                next_page_token = rsp.get("nextPageToken")
+                items: list[dict[str, Any]] = rsp.get(key)
+
+                if items is None or len(items) == 0:
+                    looping = False
+                    break
+
+                if rsp.get("nextPageToken") is None:
+                    looping = False
+                else:
+                    next_page_token = rsp.get("nextPageToken")
+
+                for item in items:
+                    yield item
+        except Exception as e:
+            logger.error(f"There was an error listing {key} with {args}: {str(e)}")
+            raise e
+
+    def load_identity_groups(self, logger: logging.Logger) -> Dict[str, Any]:
+        domain = self.config.delivery_method.domain
+        if not domain:
+            logger.info("No domain provided. Trying to fetch identities from the user workspace.")
+            api_args = {"customer": "my_customer"}
+        else:
+            api_args = {"domain": domain}
+
+        users_api = self.google_directory_service.users()
+        groups_api = self.google_directory_service.groups()
+        members_api = self.google_directory_service.members()
+
+        for user in self._get_looping_google_api_list_response(users_api, "users", args=api_args, logger=logger):
+            rfp = RemoteIdentity(
+                id=uuid.uuid4(),
+                remote_id=user["primaryEmail"],
+                name=user["name"]["fullName"] if user["name"] is not None else None,
+                email_address=user["primaryEmail"],
+                member_email_addresses=[x["address"] for x in user["emails"]],
+                type=RemoteIdentityType.USER,
+                modified_at=datetime_now(),
+            )
+            yield rfp.dict()
+
+        for group in self._get_looping_google_api_list_response(groups_api, "groups", args=api_args, logger=logger):
+            rfp = RemoteIdentity(
+                id=uuid.uuid4(),
+                remote_id=group["email"],
+                name=group["name"],
+                email_address=group["email"],
+                type=RemoteIdentityType.GROUP,
+                modified_at=datetime_now(),
+            )
+
+            for member in self._get_looping_google_api_list_response(members_api, "members", {"groupKey": group["id"]}, logger):
+                rfp.member_email_addresses = rfp.member_email_addresses or []
+                rfp.member_email_addresses.append(member["email"])
+
+            yield rfp.dict()
+
+    @property
+    def file_permissions_schema(self) -> Dict[str, Any]:
+        return ResourceSchemaLoader(package_name_from_class(self.__class__)).get_schema("file_permissions")
+
+    @property
+    def identities_schema(self) -> Dict[str, Any]:
+        return ResourceSchemaLoader(package_name_from_class(self.__class__)).get_schema("identities")

{airbyte_source_google_drive-0.1.1 → airbyte_source_google_drive-0.2.0}/source_google_drive/utils.py

@@ -15,6 +15,8 @@ def get_folder_id(url_string: str) -> str:
         if parsed_url.scheme != "https" or parsed_url.netloc != "drive.google.com":
             raise ValueError("Folder URL has to be of the form https://drive.google.com/drive/folders/<folder_id>")
         path_segments = list(filter(None, parsed_url.path.split("/")))
+        if path_segments[-1] in ["my-drive", "home"]:
+            return "root"
         if path_segments[-2] != "folders" or len(path_segments) < 3:
             raise ValueError("Folder URL has to be of the form https://drive.google.com/drive/folders/<folder_id>")
         return path_segments[-1]