airbyte-source-google-drive 0.2.0.dev202501241824__tar.gz → 0.2.0.dev202501270320__tar.gz

{airbyte_source_google_drive-0.2.0.dev202501241824 → airbyte_source_google_drive-0.2.0.dev202501270320}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: airbyte-source-google-drive
-Version: 0.2.0.dev202501241824
+Version: 0.2.0.dev202501270320
 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.26.0.dev04101)
+Requires-Dist: airbyte-cdk[file-based] (==6.26.0.dev04105)
 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.2.0.dev202501241824 → airbyte_source_google_drive-0.2.0.dev202501270320}/pyproject.toml

@@ -5,7 +5,7 @@ requires = [
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
-version = "0.2.0.dev202501241824"
+version = "0.2.0.dev202501270320"
 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.26.0.dev04101"
+version = "6.26.0.dev04105"
 
 [tool.poetry.scripts]
 source-google-drive = "source_google_drive.run:run"

{airbyte_source_google_drive-0.2.0.dev202501241824 → airbyte_source_google_drive-0.2.0.dev202501270320}/source_google_drive/spec.py

@@ -2,9 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import uuid
-from datetime import datetime
-from typing import Any, Dict, Literal, Union
+from typing import Any, Dict, Literal, Optional, Union
 
 import dpath.util
 from pydantic.v1 import BaseModel, Field
@@ -19,34 +17,15 @@ from airbyte_cdk.sources.file_based.config.abstract_file_based_spec import (
 )
 
 
-class RemoteFileIdentity(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: str
-    modified_at: datetime
-
-
-class RemoteFileMetadata(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 DeliverRecords(DeliverRecordsBase):
     # Overriding to make visible with airbyte_hidden=False
-    sync_metadata: bool = Field(
-        title="Make stream sync files metadata",
-        description="If enabled, streams will sync files metadata instead of files data.",
+    sync_acl_permissions: bool = Field(
+        title="Include ACL Permissions",
+        description="Joins Document allowlists to each stream.",
         default=False,
         airbyte_hidden=False,
     )
+    domain: Optional[str] = Field(title="Domain", description="The domain of the identities.", airbyte_hidden=False)
 
 
 class OAuthCredentials(BaseModel):

{airbyte_source_google_drive-0.2.0.dev202501241824 → airbyte_source_google_drive-0.2.0.dev202501270320}/source_google_drive/stream_reader.py

@@ -10,20 +10,25 @@ import uuid
 from datetime import datetime
 from io import IOBase
 from os.path import getsize
-from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
+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
 
 from airbyte_cdk import AirbyteTracedException, FailureType
+from airbyte_cdk.sources.file_based.config.permissions import RemoteFileIdentity, RemoteFileIdentityType, RemoteFilePermissions
 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 source_google_drive.utils import get_folder_id
 
 from .exceptions import ErrorDownloadingFile, ErrorFetchingMetadata
-from .spec import RemoteFileIdentity, RemoteFileMetadata, SourceGoogleDriveSpec
+from .spec import SourceGoogleDriveSpec
+
+# remove this, just using while test credentials are prepared
+from .temp_mock import _get_looping_google_api_list_response
 
 
 FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
@@ -73,6 +78,10 @@ PERMISSIONS_API_SCOPES = [
 ]
 
 
+def datetime_now() -> datetime:
+    return datetime.now(pytz.UTC)
+
+
 class GoogleDriveRemoteFile(RemoteFile):
     id: str
     # The mime type of the file as returned by the Google Drive API
@@ -116,7 +125,7 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
                 if self.config.credentials.auth_type == "Client":
                     creds = credentials.Credentials.from_authorized_user_info(self.config.credentials.dict())
                 else:
-                    scopes = PERMISSIONS_API_SCOPES if self.sync_metadata() else None
+                    scopes = PERMISSIONS_API_SCOPES if self.sync_acl_permissions() else None
                     creds = service_account.Credentials.from_service_account_info(
                         json.loads(self.config.credentials.service_account_info), scopes=scopes
                     )
@@ -131,6 +140,76 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
 
         return self._drive_service
 
+    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.info(f"There was an error listing {key} with {args}: {str(e)}")
+            logger.info(f"Backing off to mocked data for development purposes")
+            yield from _get_looping_google_api_list_response(service, key, args)
+
+    def load_identity_groups(self, logger: logging.Logger) -> Dict[str, Any]:
+        domain = self.config.delivery_method.domain
+        if not domain:
+            raise Exception("No domain was provided")
+        if self.google_drive_service is None or self.google_drive_service._http.credentials is None:
+            raise Exception("No auth found")
+
+        directory_service = build("admin", "directory_v1", credentials=self.google_drive_service._http.credentials)
+        users_api = directory_service.users()
+        groups_api = directory_service.groups()
+        members_api = directory_service.members()
+
+        # here is failing
+        for user in self._get_looping_google_api_list_response(users_api, "users", {"domain": domain}, logger):
+            rfp = RemoteFileIdentity(
+                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=RemoteFileIdentityType.USER,
+                modified_at=datetime_now(),
+            )
+            yield rfp.dict()
+
+        for group in self._get_looping_google_api_list_response(groups_api, "groups", {"domain": domain}, logger):
+            rfp = RemoteFileIdentity(
+                id=uuid.uuid4(),
+                remote_id=group["email"],
+                name=group["name"],
+                email_address=group["email"],
+                type=RemoteFileIdentityType.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()
+
+        return ""
+
     def get_matching_files(self, globs: List[str], prefix: Optional[str], logger: logging.Logger) -> Iterable[RemoteFile]:
         """
         Get all files matching the specified glob patterns.
@@ -356,9 +435,9 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
             description=None,
         )
 
-    def get_file_metadata(self, file: GoogleDriveRemoteFile, logger: logging.Logger) -> Dict[str, Any]:
+    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 RemoteFileMetadata(
+        return RemoteFilePermissions(
             id=file.id,
             file_path=file.uri,
             allowed_identity_remote_ids=[p.remote_id for p in remote_identities],

airbyte_source_google_drive-0.2.0.dev202501270320/source_google_drive/temp_mock.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from typing import Any, Iterator
+
+
+def _get_looping_google_api_list_response(service: Any, key: str, args: dict[str, Any]) -> Iterator[dict[str, Any]]:
+    domain = args.get("domain") or args.get("groupKey")
+
+    if key == "users":
+        # Mocking 5 users
+        for i in range(1, 6):
+            yield {
+                "primaryEmail": f"user{i}@{domain}",
+                "name": {"fullName": f"User {i}"},
+                "emails": [{"address": f"user{i}@{domain}"}],
+            }
+
+    elif key == "groups":
+        # Mocking 5 groups
+        for i in range(1, 6):
+            yield {
+                "email": f"group{i}@{domain}",
+                "name": f"Group {i}",
+                "id": f"group-id-{i}",
+            }
+
+    elif key == "members":
+        # Mocking 5 members per group
+        group_id = args.get("groupKey")
+        for i in range(1, 6):
+            yield {
+                "email": f"member{i}@{domain}",
+                "role": "MEMBER",
+                "type": "USER",
+                "groupId": group_id,
+            }