1
+ #
2
+ # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
3
+ #
4
+
1
5
import json
6
+ import logging
7
+ from dataclasses import InitVar , dataclass , field
2
8
from pathlib import Path
3
- from typing import Optional
9
+ from typing import Optional , Mapping , Union , Any
4
10
11
+ from airbyte_cdk .sources .declarative .interpolation .interpolated_string import (
12
+ InterpolatedString ,
13
+ )
5
14
from airbyte_cdk .models import AirbyteRecordMessageFileReference
6
15
from airbyte_cdk .sources .declarative .extractors .record_extractor import RecordExtractor
7
16
from airbyte_cdk .sources .declarative .partition_routers .substream_partition_router import (
8
17
SafeResponse ,
9
18
)
10
19
from airbyte_cdk .sources .declarative .requesters import Requester
11
20
from airbyte_cdk .sources .declarative .types import Record , StreamSlice
21
+ from airbyte_cdk .sources .types import Config
12
22
from airbyte_cdk .sources .utils .files_directory import get_files_directory
13
23
24
+ logger = logging .getLogger ("airbyte" )
14
25
26
+
27
+ @dataclass
15
28
class FileUploader :
16
- def __init__ (
17
- self ,
18
- requester : Requester ,
19
- download_target_extractor : RecordExtractor ,
20
- content_extractor : Optional [RecordExtractor ] = None ,
21
- ) -> None :
22
- self ._requester = requester
23
- self ._download_target_extractor = download_target_extractor
24
- self ._content_extractor = content_extractor
29
+ requester : Requester
30
+ download_target_extractor : RecordExtractor
31
+ config : Config
32
+ parameters : InitVar [Mapping [str , Any ]]
33
+
34
+ filename_extractor : Union [InterpolatedString , str ]
35
+ content_extractor : Optional [RecordExtractor ] = None
36
+
37
+ def __post_init__ (self , parameters : Mapping [str , Any ]) -> None :
38
+ self ._filename_extractor = InterpolatedString .create (
39
+ self .filename_extractor ,
40
+ parameters = parameters ,
41
+ )
25
42
26
43
def upload (self , record : Record ) -> None :
27
44
mocked_response = SafeResponse ()
28
45
mocked_response .content = json .dumps (record .data ).encode ("utf-8" )
29
- download_target = list (self ._download_target_extractor .extract_records (mocked_response ))[0 ]
46
+ download_target = list (self .download_target_extractor .extract_records (mocked_response ))[0 ]
30
47
if not isinstance (download_target , str ):
31
48
raise ValueError (
32
49
f"download_target is expected to be a str but was { type (download_target )} : { download_target } "
33
50
)
34
51
35
- response = self ._requester .send_request (
52
+ response = self .requester .send_request (
36
53
stream_slice = StreamSlice (
37
54
partition = {}, cursor_slice = {}, extra_fields = {"download_target" : download_target }
38
55
),
39
56
)
40
57
41
- if self ._content_extractor :
58
+ if self .content_extractor :
42
59
raise NotImplementedError ("TODO" )
43
60
else :
44
61
files_directory = Path (get_files_directory ())
45
- # TODO:: we could either interpolate record data if some relative_path is provided or
46
- # use partition_field value in the slice {"partition_field": some_value_id} to create a path
47
- file_relative_path = Path (record .stream_name ) / record .data ["file_name" ]
62
+
63
+ relative_path = self ._filename_extractor .eval (self .config , record = record )
64
+ relative_path = relative_path .lstrip ("/" )
65
+ file_relative_path = Path (relative_path )
48
66
49
67
full_path = files_directory / file_relative_path
50
68
full_path .parent .mkdir (parents = True , exist_ok = True )
@@ -53,8 +71,13 @@ def upload(self, record: Record) -> None:
53
71
f .write (response .content )
54
72
file_size_bytes = full_path .stat ().st_size
55
73
74
+ logger .info ("File uploaded successfully" )
75
+ logger .info (f"File url: { str (full_path )} " )
76
+ logger .info (f"File size: { file_size_bytes / 1024 } KB" )
77
+ logger .info (f"File relative path: { str (file_relative_path )} " )
78
+
56
79
record .file_reference = AirbyteRecordMessageFileReference (
57
- file_url = download_target ,
80
+ file_url = str ( full_path ) ,
58
81
file_relative_path = str (file_relative_path ),
59
82
file_size_bytes = file_size_bytes ,
60
83
)
0 commit comments