Merge branch 'aj/feat/add-standard-tests-cli' into aj/feat/add-cli-build-command

Author: aaronsteers

.github/workflows/connector-tests.yml

@@ -75,11 +75,11 @@ jobs:
           # Chargebee is being flaky:
           # - connector: source-chargebee
           #   cdk_extra: n/a
-          # These two are behind in CDK updates and can't be used as tests until they are updated:
-          # - connector: source-s3
-          #   cdk_extra: file-based
+          # This one is behind in CDK updates and can't be used as tests until it is updated:
           # - connector: destination-pinecone
           #   cdk_extra: vector-db-based
+          - connector: source-google-drive
+            cdk_extra: file-based
           - connector: destination-motherduck
             cdk_extra: sql
           # ZenDesk currently failing (as of 2024-12-02)

.github/workflows/poe-command.yml

@@ -0,0 +1,28 @@
+name: On-Demand Poe Task
+
+on:
+  workflow_dispatch:
+    inputs:
+      comment-id:
+        description: "Optional comment-id of the slash command. Ignore if not applicable."
+        required: false
+      pr:
+        description: "PR Number"
+        required: false
+
+permissions:
+  contents: write
+  pull-requests: write
+
+jobs:
+  run-poe-command:
+    env:
+      GCP_GSM_CREDENTIALS: ${{ secrets.GCP_GSM_CREDENTIALS }}
+    runs-on: ubuntu-latest
+    steps:
+      - name: Run Poe Slash Command Processor
+        uses: aaronsteers/poe-command-processor@v1
+        with:
+          pr: ${{ github.event.inputs.pr }}
+          comment-id: ${{ github.event.inputs.comment-id }}
+          github-token: ${{ secrets.GH_PAT_MAINTENANCE_OCTAVIA }}

.github/workflows/slash_command_dispatch.yml

@@ -6,8 +6,6 @@ on:
 
 jobs:
   slashCommandDispatch:
-    # Only allow slash commands on pull request (not on issues)
-    if: ${{ github.event.issue.pull_request }}
     runs-on: ubuntu-24.04
     steps:
       - name: Slash Command Dispatch
@@ -19,17 +17,25 @@ jobs:
           repository: ${{ github.repository }}
           token: ${{ secrets.GH_PAT_MAINTENANCE_OCTAVIA }}
           dispatch-type: workflow
-          issue-type: pull-request
+          issue-type: both
+
+          # Only run for users with 'write' permission on the main repository
+          permission: write
+
           commands: |
             autofix
             test
             poetry-lock
+            poe
+
+          # Notes regarding static-args:
+          # - Slash commands can be invoked from both issues and comments.
+          # - If the slash command is invoked from an issue, we intentionally pass 'null' as the PR number.
+          # - Comment ID will always be sent, and this is sufficient to post back status updates to the originating comment.
           static-args: |
-            pr=${{ github.event.issue.number }}
+            pr=${{ github.event.issue.pull_request != null && github.event.issue.number || '' }}
             comment-id=${{ github.event.comment.id }}
 
-          # Only run for users with 'write' permission on the main repository
-          permission: write
 
       - name: Edit comment with error message
         if: steps.dispatch.outputs.error-message

airbyte_cdk/models/__init__.py

@@ -19,6 +19,7 @@
     AirbyteMessage,
     AirbyteProtocol,
     AirbyteRecordMessage,
+    AirbyteRecordMessageFileReference,
     AirbyteStateBlob,
     AirbyteStateMessage,
     AirbyteStateStats,

airbyte_cdk/models/airbyte_protocol.py

@@ -8,8 +8,6 @@
 from airbyte_protocol_dataclasses.models import *  # noqa: F403  # Allow '*'
 from serpyco_rs.metadata import Alias
 
-from airbyte_cdk.models.file_transfer_record_message import AirbyteFileTransferRecordMessage
-
 # ruff: noqa: F405  # ignore fuzzy import issues with 'import *'
 
 
@@ -84,7 +82,7 @@ class AirbyteMessage:
     spec: Optional[ConnectorSpecification] = None  # type: ignore [name-defined]
     connectionStatus: Optional[AirbyteConnectionStatus] = None  # type: ignore [name-defined]
     catalog: Optional[AirbyteCatalog] = None  # type: ignore [name-defined]
-    record: Optional[Union[AirbyteFileTransferRecordMessage, AirbyteRecordMessage]] = None  # type: ignore [name-defined]
+    record: Optional[AirbyteRecordMessage] = None  # type: ignore [name-defined]
     state: Optional[AirbyteStateMessage] = None
     trace: Optional[AirbyteTraceMessage] = None  # type: ignore [name-defined]
     control: Optional[AirbyteControlMessage] = None  # type: ignore [name-defined]

airbyte_cdk/models/file_transfer_record_message.py

@@ -149,7 +149,7 @@ def on_record(self, record: Record) -> Iterable[AirbyteMessage]:
         message = stream_data_to_airbyte_message(
             stream_name=record.stream_name,
             data_or_message=record.data,
-            is_file_transfer_message=record.is_file_transfer_message,
+            file_reference=record.file_reference,
         )
         stream = self._stream_name_to_instance[record.stream_name]
 

airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -28,6 +28,7 @@
     PerPartitionWithGlobalCursor,
 )
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
+from airbyte_cdk.sources.declarative.models import FileUploader
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ConcurrencyLevel as ConcurrencyLevelModel,
 )
@@ -209,6 +210,11 @@ def _group_streams(
             # these legacy Python streams the way we do low-code streams to determine if they are concurrent compatible,
             # so we need to treat them as synchronous
 
+            supports_file_transfer = (
+                isinstance(declarative_stream, DeclarativeStream)
+                and "file_uploader" in name_to_stream_mapping[declarative_stream.name]
+            )
+
             if (
                 isinstance(declarative_stream, DeclarativeStream)
                 and name_to_stream_mapping[declarative_stream.name]["type"]
@@ -325,6 +331,7 @@ def _group_streams(
                             else None,
                             logger=self.logger,
                             cursor=cursor,
+                            supports_file_transfer=supports_file_transfer,
                         )
                     )
                 elif (
@@ -356,6 +363,7 @@ def _group_streams(
                             cursor_field=None,
                             logger=self.logger,
                             cursor=final_state_cursor,
+                            supports_file_transfer=supports_file_transfer,
                         )
                     )
                 elif (
@@ -410,6 +418,7 @@ def _group_streams(
                             cursor_field=perpartition_cursor.cursor_field.cursor_field_key,
                             logger=self.logger,
                             cursor=perpartition_cursor,
+                            supports_file_transfer=supports_file_transfer,
                         )
                     )
                 else:

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -1448,6 +1448,42 @@ definitions:
             - "$ref": "#/definitions/LegacyToPerPartitionStateMigration"
             - "$ref": "#/definitions/CustomStateMigration"
         default: []
+      file_uploader:
+        title: File Uploader
+        description: (experimental) Describes how to fetch a file
+        type: object
+        required:
+          - type
+          - requester
+          - download_target_extractor
+        properties:
+          type:
+            type: string
+            enum: [ FileUploader ]
+          requester:
+            description: Requester component that describes how to prepare HTTP requests to send to the source API.
+            anyOf:
+              - "$ref": "#/definitions/CustomRequester"
+              - "$ref": "#/definitions/HttpRequester"
+          download_target_extractor:
+            description: Responsible for fetching the url where the file is located. This is applied on each records and not on the HTTP response
+            anyOf:
+              - "$ref": "#/definitions/CustomRecordExtractor"
+              - "$ref": "#/definitions/DpathExtractor"
+          file_extractor:
+            description: Responsible for fetching the content of the file. If not defined, the assumption is that the whole response body is the file content
+            anyOf:
+              - "$ref": "#/definitions/CustomRecordExtractor"
+              - "$ref": "#/definitions/DpathExtractor"
+          filename_extractor:
+            description: Defines the name to store the file. Stream name is automatically added to the file path. File unique ID can be used to avoid overwriting files. Random UUID will be used if the extractor is not provided.
+            type: string
+            interpolation_context:
+              - config
+              - record
+            examples:
+              - "{{ record.id }}/{{ record.file_name }}/"
+              - "{{ record.id }}_{{ record.file_name }}/"
       $parameters:
         type: object
         additional_properties: true

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -15,6 +15,7 @@
 )
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.models import SchemaNormalization
+from airbyte_cdk.sources.declarative.retrievers.file_uploader import FileUploader
 from airbyte_cdk.sources.declarative.transformations import RecordTransformation
 from airbyte_cdk.sources.types import Config, Record, StreamSlice, StreamState
 from airbyte_cdk.sources.utils.transform import TypeTransformer
@@ -42,6 +43,7 @@ class RecordSelector(HttpSelector):
     record_filter: Optional[RecordFilter] = None
     transformations: List[RecordTransformation] = field(default_factory=lambda: [])
     transform_before_filtering: bool = False
+    file_uploader: Optional[FileUploader] = None
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._parameters = parameters
@@ -117,7 +119,10 @@ def filter_and_transform(
             transformed_filtered_data, schema=records_schema
         )
         for data in normalized_data:
-            yield Record(data=data, stream_name=self.name, associated_slice=stream_slice)
+            record = Record(data=data, stream_name=self.name, associated_slice=stream_slice)
+            if self.file_uploader:
+                self.file_uploader.upload(record)
+            yield record
 
     def _normalize_by_schema(
         self, records: Iterable[Mapping[str, Any]], schema: Optional[Mapping[str, Any]]