Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion airbyte_cdk/sources/file_based/file_based_stream_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import logging
import time
import uuid
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from enum import Enum
Expand Down Expand Up @@ -254,6 +255,11 @@ def _get_file_transfer_paths(
- LOCAL_FILE_PATH: The absolute path to the file.
- FILE_NAME: The name of the referenced file.
- FILE_FOLDER: The folder of the referenced file.

The local path is namespaced under a unique staging subdirectory so that two source files
resolving to the same relative path never share a staging file. Sharing a staging file makes
one download overwrite the other and makes the destination fail once it has consumed and
deleted the first copy.
"""
preserve_directory_structure = self.preserve_directory_structure()

Expand All @@ -264,7 +270,7 @@ def _get_file_transfer_paths(
file_relative_path = source_file_relative_path.lstrip("/")
else:
file_relative_path = file_name
local_file_path = path.join(staging_directory, file_relative_path)
local_file_path = path.join(staging_directory, uuid.uuid4().hex, file_relative_path)
# Ensure the local directory exists
makedirs(path.dirname(local_file_path), exist_ok=True)

Expand Down
66 changes: 59 additions & 7 deletions unit_tests/sources/file_based/test_file_based_stream_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -433,7 +433,7 @@ def test_globs_and_prefixes_from_globs(


@pytest.mark.parametrize(
"config, source_file_path, expected_file_relative_path, expected_local_file_path",
"config, source_file_path, expected_file_relative_path",
[
pytest.param(
{
Expand All @@ -445,7 +445,6 @@ def test_globs_and_prefixes_from_globs(
},
"mirror_paths_testing/not_duplicates/data/jan/monthly-kickoff-202402.mpeg",
"mirror_paths_testing/not_duplicates/data/jan/monthly-kickoff-202402.mpeg",
f"{files_directory}/mirror_paths_testing/not_duplicates/data/jan/monthly-kickoff-202402.mpeg",
id="preserve_directories_present_and_true",
),
pytest.param(
Expand All @@ -458,21 +457,18 @@ def test_globs_and_prefixes_from_globs(
},
"mirror_paths_testing/not_duplicates/data/jan/monthly-kickoff-202402.mpeg",
"monthly-kickoff-202402.mpeg",
f"{files_directory}/monthly-kickoff-202402.mpeg",
id="preserve_directories_present_and_false",
),
pytest.param(
{"streams": [], "delivery_method": {"delivery_type": "use_file_transfer"}},
"mirror_paths_testing/not_duplicates/data/jan/monthly-kickoff-202402.mpeg",
"mirror_paths_testing/not_duplicates/data/jan/monthly-kickoff-202402.mpeg",
f"{files_directory}/mirror_paths_testing/not_duplicates/data/jan/monthly-kickoff-202402.mpeg",
id="preserve_directories_not_present_defaults_true",
),
pytest.param(
{"streams": []},
"mirror_paths_testing/not_duplicates/data/jan/monthly-kickoff-202402.mpeg",
"mirror_paths_testing/not_duplicates/data/jan/monthly-kickoff-202402.mpeg",
f"{files_directory}/mirror_paths_testing/not_duplicates/data/jan/monthly-kickoff-202402.mpeg",
id="file_transfer_flag_not_present_defaults_true",
),
],
Expand All @@ -481,7 +477,6 @@ def test_preserve_sub_directories_scenarios(
config: Mapping[str, Any],
source_file_path: str,
expected_file_relative_path: str,
expected_local_file_path: str,
) -> None:
"""
Test scenarios when preserve_directory_structure is True or False, the flag indicates whether we need to
Expand All @@ -496,11 +491,68 @@ def test_preserve_sub_directories_scenarios(
assert (
file_paths[AbstractFileBasedStreamReader.FILE_RELATIVE_PATH] == expected_file_relative_path
)
assert file_paths[AbstractFileBasedStreamReader.LOCAL_FILE_PATH] == expected_local_file_path
local_file_path = file_paths[AbstractFileBasedStreamReader.LOCAL_FILE_PATH]
staging_subdirectory = path.relpath(local_file_path, files_directory).split(path.sep)[0]
assert local_file_path == path.join(
files_directory, staging_subdirectory, expected_file_relative_path
)
assert file_paths[AbstractFileBasedStreamReader.FILE_NAME] == path.basename(source_file_path)
assert file_paths[AbstractFileBasedStreamReader.FILE_FOLDER] == path.dirname(source_file_path)


@pytest.mark.parametrize(
"preserve_directory_structure, first_source_file_path, second_source_file_path",
[
pytest.param(
True,
"folder_a/Untitled document.docx",
"folder_a/Untitled document.docx",
id="same_relative_path_from_distinct_source_files",
),
pytest.param(
False,
"folder_a/Untitled document.docx",
"folder_b/Untitled document.docx",
id="same_file_name_in_different_folders_without_preserved_directories",
),
],
)
def test_staging_paths_are_unique_per_file(
preserve_directory_structure: bool,
first_source_file_path: str,
second_source_file_path: str,
) -> None:
"""
Two source files that resolve to the same relative path must not share a staging file: the second
download would overwrite the first, and the destination fails with a FileNotFoundException once it
has consumed and deleted the first copy.
"""
reader = TestStreamReader()
reader.config = TestSpec(
streams=[],
delivery_method={
"delivery_type": "use_file_transfer",
"preserve_directory_structure": preserve_directory_structure,
},
)

first_paths = reader._get_file_transfer_paths(
first_source_file_path, staging_directory=f"{files_directory}/"
)
second_paths = reader._get_file_transfer_paths(
second_source_file_path, staging_directory=f"{files_directory}/"
)

assert (
first_paths[AbstractFileBasedStreamReader.FILE_RELATIVE_PATH]
== second_paths[AbstractFileBasedStreamReader.FILE_RELATIVE_PATH]
)
assert (
first_paths[AbstractFileBasedStreamReader.LOCAL_FILE_PATH]
!= second_paths[AbstractFileBasedStreamReader.LOCAL_FILE_PATH]
)


@pytest.mark.parametrize(
"start_date_str, expected",
[
Expand Down
Loading