airbyte-source-google-drive 0.2.0.dev202502032236__py3-none-any.whl → 0.2.0.dev202502061746__py3-none-any.whl

{airbyte_source_google_drive-0.2.0.dev202502032236.dist-info → airbyte_source_google_drive-0.2.0.dev202502061746.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: airbyte-source-google-drive
-Version: 0.2.0.dev202502032236
+Version: 0.2.0.dev202502061746
 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.dev04107)
+Requires-Dist: airbyte-cdk[file-based] (==6.26.0.dev04108)
 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.dev202502032236.dist-info → airbyte_source_google_drive-0.2.0.dev202502061746.dist-info}/RECORD

@@ -1,11 +1,11 @@
 source_google_drive/__init__.py,sha256=SAgPIoGpjdHTY7mvf__A99FiElBh3WvdKSHLstg-BUE,134
 source_google_drive/exceptions.py,sha256=6sBxbdZhj0CARkctfBaNgGHB_GGcFUkY4cnMUWyOruM,268
 source_google_drive/run.py,sha256=5AQd906FbSGaNZj0udrr-VKmiSS4nsor7F7noMjsdi0,677
-source_google_drive/source.py,sha256=OyKTjPg80KR8yZiYiC3YSi6E6egothtAn3YAHAS6Q6A,2839
+source_google_drive/source.py,sha256=tOljkgeg8RIxmhrpfa-qgGhI6WPhiL_-I8itTa6JH04,3959
 source_google_drive/spec.py,sha256=-WkA2zGuQtf3G7uK8uq9BnimUlQh0s3vsqROmIHOgzI,4718
-source_google_drive/stream_reader.py,sha256=8F0jEKz1ElM2wqHxVR91RUJfCSDoLtMQUxbX5lFkEy0,20166
+source_google_drive/stream_reader.py,sha256=jKMibQGjpP-9raJ557TqbOg6iGVuPYv9vamStwXDOCQ,20176
 source_google_drive/utils.py,sha256=ewR-kBKLmtD-s7zqCfGECfzWYF43tpQdscAQIlUEkR8,1022
-airbyte_source_google_drive-0.2.0.dev202502032236.dist-info/METADATA,sha256=Hfb2iDOlay912TAmyj3QZ9GqvyVstYZflM_J2lpf2Lk,5536
-airbyte_source_google_drive-0.2.0.dev202502032236.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
-airbyte_source_google_drive-0.2.0.dev202502032236.dist-info/entry_points.txt,sha256=YgpJf0nA5Mn0B7YC9VOFI847vz1jI6U4q7BeLUOXa54,67
-airbyte_source_google_drive-0.2.0.dev202502032236.dist-info/RECORD,,
+airbyte_source_google_drive-0.2.0.dev202502061746.dist-info/METADATA,sha256=pqQMdjznFYJxr7fcCYHXaLBZcESfHaSS7o0jr2PUH2s,5536
+airbyte_source_google_drive-0.2.0.dev202502061746.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
+airbyte_source_google_drive-0.2.0.dev202502061746.dist-info/entry_points.txt,sha256=YgpJf0nA5Mn0B7YC9VOFI847vz1jI6U4q7BeLUOXa54,67
+airbyte_source_google_drive-0.2.0.dev202502061746.dist-info/RECORD,,

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",

source_google_drive/stream_reader.py

@@ -63,11 +63,8 @@ PUBLIC_PERMISSION_IDS = [
     "domainWithLink",
 ]
 
-PERMISSIONS_API_SCOPES = [
-    "https://www.googleapis.com/auth/drive",
-    "https://www.googleapis.com/auth/drive.readonly",
-    "https://www.googleapis.com/auth/drive.metadata.readonly",
-    "https://www.googleapis.com/auth/drive.file",
+
+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",
@@ -91,6 +88,7 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
     def __init__(self):
         super().__init__()
         self._drive_service = None
+        self._directory_service = None
 
     @property
     def config(self) -> SourceGoogleDriveSpec:
@@ -110,32 +108,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:
-                    scopes = PERMISSIONS_API_SCOPES if self.include_identities_stream() else None
-                    creds = service_account.Credentials.from_service_account_info(
-                        json.loads(self.config.credentials.service_account_info), scopes=scopes
-                    )
-                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.
@@ -340,7 +347,6 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
                 if identity is not None:
                     remote_identities.append(identity)
 
-            logger.info(f"File {file_name} has {len(remote_identities)} valid permissions")
             return remote_identities, is_public
         except Exception as e:
             raise ErrorFetchingMetadata(f"An error occurred while retrieving file permissions: {str(e)}")
@@ -399,16 +405,16 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
     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")
+            logger.info("No domain provided. Trying to fetch identities from the user workspace.")
+            api_args = {"customer": "my_customer"}
+        else:
+            api_args = {"domain": domain}
 
-        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()
+        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", {"domain": domain}, logger):
+        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"],
@@ -420,7 +426,7 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
             )
             yield rfp.dict()
 
-        for group in self._get_looping_google_api_list_response(groups_api, "groups", {"domain": domain}, logger):
+        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"],