airbyte-source-google-drive 0.0.12__tar.gz → 0.1.0__tar.gz

{airbyte_source_google_drive-0.0.12 → airbyte_source_google_drive-0.1.0}/PKG-INFO

@@ -1,23 +1,22 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: airbyte-source-google-drive
-Version: 0.0.12
+Version: 0.1.0
 Summary: Source implementation for Google Drive.
-Home-page: https://airbyte.com
 License: ELv2
 Author: Airbyte
 Author-email: contact@airbyte.io
-Requires-Python: >=3.9,<3.12
+Requires-Python: >=3.10,<3.12
 Classifier: License :: Other/Proprietary License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: airbyte-cdk[file-based] (>=1,<2)
+Requires-Dist: airbyte-cdk[file-based] (>=6.18.2,<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)
 Requires-Dist: google-auth-oauthlib (==1.1.0)
 Project-URL: Documentation, https://docs.airbyte.com/integrations/sources/google-drive
+Project-URL: Homepage, https://airbyte.com
 Project-URL: Repository, https://github.com/airbytehq/airbyte
 Description-Content-Type: text/markdown
 

{airbyte_source_google_drive-0.0.12 → airbyte_source_google_drive-0.1.0}/pyproject.toml

@@ -5,7 +5,7 @@ requires = [
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
-version = "0.0.12"
+version = "0.1.0"
 name = "airbyte-source-google-drive"
 description = "Source implementation for Google Drive."
 authors = [
@@ -21,7 +21,7 @@ packages = [
 ]
 
 [tool.poetry.dependencies]
-python = "^3.9,<3.12"
+python = "^3.10,<3.12"
 google-api-python-client = "==2.104.0"
 google-auth-httplib2 = "==0.1.1"
 google-auth-oauthlib = "==1.1.0"
@@ -31,7 +31,7 @@ google-api-python-client-stubs = "==1.18.0"
 extras = [
     "file-based",
 ]
-version = "^1"
+version = "^6.18.2"
 
 [tool.poetry.scripts]
 source-google-drive = "source_google_drive.run:run"

airbyte_source_google_drive-0.1.0/source_google_drive/exceptions.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from airbyte_cdk.sources.file_based.exceptions import BaseFileBasedSourceError
+
+
+class ErrorFetchingMetadata(BaseFileBasedSourceError):
+    pass
+
+
+class ErrorDownloadingFile(BaseFileBasedSourceError):
+    pass

{airbyte_source_google_drive-0.0.12 → airbyte_source_google_drive-0.1.0}/source_google_drive/source.py

@@ -6,6 +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.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
@@ -32,7 +33,7 @@ class SourceGoogleDrive(FileBasedSource):
             documentationUrl=self.spec_class.documentation_url(),
             connectionSpecification=self.spec_class.schema(),
             advanced_auth=AdvancedAuth(
-                auth_flow_type="oauth2.0",
+                auth_flow_type=AuthFlowType.oauth2_0,
                 predicate_key=["credentials", "auth_type"],
                 predicate_value="Client",
                 oauth_config_specification=OAuthConfigSpecification(

{airbyte_source_google_drive-0.0.12 → airbyte_source_google_drive-0.1.0}/source_google_drive/spec.py

@@ -6,9 +6,10 @@
 from typing import Any, Dict, Literal, 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
-from pydantic import BaseModel, Field
+from airbyte_cdk.sources.file_based.config.abstract_file_based_spec import AbstractFileBasedSpec, DeliverRawFiles, DeliverRecords
 
 
 class OAuthCredentials(BaseModel):
@@ -59,6 +60,16 @@ class SourceGoogleDriveSpec(AbstractFileBasedSpec, BaseModel):
         pattern_descriptor="https://drive.google.com/drive/folders/MY-FOLDER-ID",
     )
 
+    delivery_method: DeliverRecords | DeliverRawFiles = Field(
+        title="Delivery Method",
+        discriminator="delivery_type",
+        type="object",
+        order=1,
+        display_type="radio",
+        group="advanced",
+        default="use_records_transfer",
+    )
+
     credentials: Union[OAuthCredentials, ServiceAccountCredentials] = Field(
         title="Authentication", description="Credentials for connecting to the Google Drive API", discriminator="auth_type", type="object"
     )

{airbyte_source_google_drive-0.0.12 → airbyte_source_google_drive-0.1.0}/source_google_drive/stream_reader.py

@@ -8,25 +8,51 @@ import json
 import logging
 from datetime import datetime
 from io import IOBase
-from typing import Iterable, List, Optional, Set
+from os.path import getsize
+from typing import Dict, Iterable, List, Optional, Set
 
-from airbyte_cdk import AirbyteTracedException, FailureType
-from airbyte_cdk.sources.file_based.file_based_stream_reader import AbstractFileBasedStreamReader, FileReadMode
-from airbyte_cdk.sources.file_based.remote_file import RemoteFile
 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.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 SourceGoogleDriveSpec
 
+
 FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
 GOOGLE_DOC_MIME_TYPE = "application/vnd.google-apps.document"
-EXPORTABLE_DOCUMENTS_MIME_TYPES = [
-    GOOGLE_DOC_MIME_TYPE,
-    "application/vnd.google-apps.presentation",
-    "application/vnd.google-apps.drawing",
-]
+GOOGLE_PRESENTATION_MIME_TYPE = "application/vnd.google-apps.presentation"
+GOOGLE_SPREADSHEET_MIME_TYPE = "application/vnd.google-apps.spreadsheet"
+GOOGLE_DRAWING_MIME_TYPE = "application/vnd.google-apps.drawing"
+EXPORTABLE_DOCUMENTS_MIME_TYPES = [GOOGLE_DOC_MIME_TYPE, GOOGLE_PRESENTATION_MIME_TYPE, GOOGLE_DRAWING_MIME_TYPE]
+
+EXPORT_MEDIA_MIME_TYPE_DOC = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+EXPORT_MEDIA_MIME_TYPE_SPREADSHEET = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+EXPORT_MEDIA_MIME_TYPE_PRESENTATION = "application/vnd.openxmlformats-officedocument.presentationml.presentation"
+EXPORT_MEDIA_MIME_TYPE_PDF = "application/pdf"
+
+# This key is used to relate the required mimeType parameter for export_media() method
+EXPORT_MEDIA_MIME_TYPE_KEY = "exportable_mime_type"
+DOCUMENT_FILE_EXTENSION_KEY = "document_file_extension"
+
+DOWNLOADABLE_DOCUMENTS_MIME_TYPES = {
+    GOOGLE_DOC_MIME_TYPE: {EXPORT_MEDIA_MIME_TYPE_KEY: EXPORT_MEDIA_MIME_TYPE_DOC, DOCUMENT_FILE_EXTENSION_KEY: ".docx"},
+    GOOGLE_SPREADSHEET_MIME_TYPE: {
+        EXPORT_MEDIA_MIME_TYPE_KEY: EXPORT_MEDIA_MIME_TYPE_SPREADSHEET,
+        DOCUMENT_FILE_EXTENSION_KEY: ".xlsx",
+    },
+    GOOGLE_PRESENTATION_MIME_TYPE: {
+        EXPORT_MEDIA_MIME_TYPE_KEY: EXPORT_MEDIA_MIME_TYPE_PRESENTATION,
+        DOCUMENT_FILE_EXTENSION_KEY: ".pptx",
+    },
+    GOOGLE_DRAWING_MIME_TYPE: {EXPORT_MEDIA_MIME_TYPE_KEY: EXPORT_MEDIA_MIME_TYPE_PDF, DOCUMENT_FILE_EXTENSION_KEY: ".pdf"},
+}
 
 
 class GoogleDriveRemoteFile(RemoteFile):
@@ -37,6 +63,8 @@ class GoogleDriveRemoteFile(RemoteFile):
 
 
 class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
+    FILE_SIZE_LIMIT = 1_500_000_000
+
     def __init__(self):
         super().__init__()
         self._drive_service = None
@@ -146,6 +174,8 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
         """
         Returns true if the given file is a Google App document that can be exported.
         """
+        if self.use_file_transfer():
+            return mime_type in DOWNLOADABLE_DOCUMENTS_MIME_TYPES
         return mime_type in EXPORTABLE_DOCUMENTS_MIME_TYPES
 
     def open_file(self, file: GoogleDriveRemoteFile, mode: FileReadMode, encoding: Optional[str], logger: logging.Logger) -> IOBase:
@@ -176,10 +206,79 @@ class SourceGoogleDriveStreamReader(AbstractFileBasedStreamReader):
     def _get_export_mime_type(self, original_mime_type: str):
         """
         Returns the mime type to export Google App documents as.
-
-        Google Docs are exported as Docx to preserve as much formatting as possible, everything else goes through PDF.
         """
+        if self.use_file_transfer():
+            return DOWNLOADABLE_DOCUMENTS_MIME_TYPES[original_mime_type][EXPORT_MEDIA_MIME_TYPE_KEY]
+
         if original_mime_type.startswith(GOOGLE_DOC_MIME_TYPE):
+            # Google Docs are exported as Docx to preserve as much formatting as possible, everything else goes through PDF.
             return "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
         else:
             return "application/pdf"
+
+    def file_size(self, file: GoogleDriveRemoteFile) -> int:
+        """
+        Retrieves the size of a file in Google Drive.
+
+        Args:
+            file (RemoteFile): The file to get the size for.
+
+        Returns:
+            int: The file size in bytes.
+        ref: https://developers.google.com/drive/api/reference/rest/v3/files/get
+        """
+        try:
+            file_metadata = self.google_drive_service.files().get(fileId=file.id, fields="size", supportsAllDrives=True).execute()
+            return int(file_metadata["size"])
+        except KeyError:
+            raise ErrorFetchingMetadata(f"Size was expected in metadata response but was missing")
+        except Exception as e:
+            raise ErrorFetchingMetadata(f"An error occurred while retrieving file size: {str(e)}")
+
+    def get_file(self, file: GoogleDriveRemoteFile, local_directory: str, logger: logging.Logger) -> Dict[str, str | int]:
+        """
+        Downloads a file from Google Drive to a specified local directory.
+
+        Args:
+            file (RemoteFile): The file to download, containing its Google Drive ID.
+            local_directory (str): The local directory to save the file.
+            logger (logging.Logger): Logger for debugging and information.
+
+        Returns:
+            Dict[str, str | int]: Contains the local file path and file size in bytes.
+        ref: https://developers.google.com/drive/api/guides/manage-downloads#python
+        """
+        file_size = self.file_size(file)
+        # I'm putting this check here so we can remove the safety wheels per connector when ready.
+        if file_size > self.FILE_SIZE_LIMIT:
+            message = "File size exceeds the 1 GB limit."
+            raise FileSizeLimitError(message=message, internal_message=message, failure_type=FailureType.config_error)
+
+        try:
+            file_relative_path, local_file_path, absolute_file_path = self._get_file_transfer_paths(file, local_directory)
+
+            if self._is_exportable_document(file.original_mime_type):
+                request = self.google_drive_service.files().export_media(fileId=file.id, mimeType=file.mime_type)
+
+                file_extension = DOWNLOADABLE_DOCUMENTS_MIME_TYPES[file.original_mime_type][DOCUMENT_FILE_EXTENSION_KEY]
+                local_file_path += file_extension
+                absolute_file_path += file_extension
+                file_relative_path += file_extension
+            else:
+                request = self.google_drive_service.files().get_media(fileId=file.id)
+
+            with open(local_file_path, "wb") as local_file:
+                downloader = MediaIoBaseDownload(local_file, request)
+                done = False
+                while not done:
+                    status, done = downloader.next_chunk(num_retries=5)
+                    progress = status.resumable_progress / status.total_size * 100 if status.total_size else 0
+                    logger.info(f"Processing file {file.uri}, progress: {progress:.2f}%")
+
+            # native google objects seems to be reporting lower size through the api than the final download size
+            file_size = getsize(local_file_path)
+
+            return {"file_url": absolute_file_path, "bytes": file_size, "file_relative_path": file_relative_path}
+
+        except Exception as e:
+            raise ErrorDownloadingFile(f"There was an error while trying to download the file {file.uri}: {str(e)}")