airbyte_cdk.sources.file_based

 48class AbstractFileBasedSpec(BaseModel):
 49    """
 50    Used during spec; allows the developer to configure the cloud provider specific options
 51    that are needed when users configure a file-based source.
 52    """
 53
 54    start_date: Optional[str] = Field(
 55        title="Start Date",
 56        description="UTC date and time in the format 2017-01-25T00:00:00.000000Z. Any file modified before this date will not be replicated.",
 57        examples=["2021-01-01T00:00:00.000000Z"],
 58        format="date-time",
 59        pattern="^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}.[0-9]{6}Z$",
 60        pattern_descriptor="YYYY-MM-DDTHH:mm:ss.SSSSSSZ",
 61        order=1,
 62    )
 63
 64    streams: List[FileBasedStreamConfig] = Field(
 65        title="The list of streams to sync",
 66        description='Each instance of this configuration defines a <a href="https://docs.airbyte.com/cloud/core-concepts#stream">stream</a>. Use this to define which files belong in the stream, their format, and how they should be parsed and validated. When sending data to warehouse destination such as Snowflake or BigQuery, each stream is a separate table.',
 67        order=10,
 68    )
 69
 70    delivery_method: Union[DeliverRecords, DeliverRawFiles, DeliverPermissions] = Field(
 71        title="Delivery Method",
 72        discriminator="delivery_type",
 73        type="object",
 74        order=7,
 75        display_type="radio",
 76        group="advanced",
 77        default="use_records_transfer",
 78        airbyte_hidden=True,
 79    )
 80
 81    @classmethod
 82    @abstractmethod
 83    def documentation_url(cls) -> AnyUrl:
 84        """
 85        :return: link to docs page for this source e.g. "https://docs.airbyte.com/integrations/sources/s3"
 86        """
 87
 88    @classmethod
 89    def schema(cls, *args: Any, **kwargs: Any) -> Dict[str, Any]:
 90        """
 91        Generates the mapping comprised of the config fields
 92        """
 93        schema = super().schema(*args, **kwargs)
 94        transformed_schema: Dict[str, Any] = copy.deepcopy(schema)
 95        schema_helpers.expand_refs(transformed_schema)
 96        cls.replace_enum_allOf_and_anyOf(transformed_schema)
 97        cls.remove_discriminator(transformed_schema)
 98
 99        return transformed_schema
100
101    @staticmethod
102    def remove_discriminator(schema: Dict[str, Any]) -> None:
103        """pydantic adds "discriminator" to the schema for oneOfs, which is not treated right by the platform as we inline all references"""
104        dpath.delete(schema, "properties/**/discriminator")
105
106    @staticmethod
107    def replace_enum_allOf_and_anyOf(schema: Dict[str, Any]) -> Dict[str, Any]:
108        """
109        allOfs are not supported by the UI, but pydantic is automatically writing them for enums.
110        Unpacks the enums under allOf and moves them up a level under the enum key
111        anyOfs are also not supported by the UI, so we replace them with the similar oneOf, with the
112        additional validation that an incoming config only matches exactly one of a field's types.
113        """
114        objects_to_check = schema["properties"]["streams"]["items"]["properties"]["format"]
115        objects_to_check["type"] = "object"
116        objects_to_check["oneOf"] = objects_to_check.pop("anyOf", [])
117        for format in objects_to_check["oneOf"]:
118            for key in format["properties"]:
119                object_property = format["properties"][key]
120                AbstractFileBasedSpec.move_enum_to_root(object_property)
121
122        properties_to_change = ["validation_policy"]
123        for property_to_change in properties_to_change:
124            property_object = schema["properties"]["streams"]["items"]["properties"][
125                property_to_change
126            ]
127            if "anyOf" in property_object:
128                schema["properties"]["streams"]["items"]["properties"][property_to_change][
129                    "type"
130                ] = "object"
131                schema["properties"]["streams"]["items"]["properties"][property_to_change][
132                    "oneOf"
133                ] = property_object.pop("anyOf")
134            AbstractFileBasedSpec.move_enum_to_root(property_object)
135
136        csv_format_schemas = list(
137            filter(
138                lambda format: format["properties"]["filetype"]["default"] == "csv",
139                schema["properties"]["streams"]["items"]["properties"]["format"]["oneOf"],
140            )
141        )
142        if len(csv_format_schemas) != 1:
143            raise ValueError(f"Expecting only one CSV format but got {csv_format_schemas}")
144        csv_format_schemas[0]["properties"]["header_definition"]["oneOf"] = csv_format_schemas[0][
145            "properties"
146        ]["header_definition"].pop("anyOf", [])
147        csv_format_schemas[0]["properties"]["header_definition"]["type"] = "object"
148        return schema
149
150    @staticmethod
151    def move_enum_to_root(object_property: Dict[str, Any]) -> None:
152        if "allOf" in object_property and "enum" in object_property["allOf"][0]:
153            object_property["enum"] = object_property["allOf"][0]["enum"]
154            object_property.pop("allOf")

 35class AbstractFileBasedStreamReader(ABC):
 36    DATE_TIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
 37    FILE_RELATIVE_PATH = "file_relative_path"
 38    FILE_NAME = "file_name"
 39    LOCAL_FILE_PATH = "local_file_path"
 40    FILE_FOLDER = "file_folder"
 41    FILE_SIZE_LIMIT = 1_500_000_000
 42
 43    def __init__(self) -> None:
 44        self._config = None
 45
 46    @property
 47    def config(self) -> Optional[AbstractFileBasedSpec]:
 48        return self._config
 49
 50    @config.setter
 51    @abstractmethod
 52    def config(self, value: AbstractFileBasedSpec) -> None:
 53        """
 54        FileBasedSource reads the config from disk and parses it, and once parsed, the source sets the config on its StreamReader.
 55
 56        Note: FileBasedSource only requires the keys defined in the abstract config, whereas concrete implementations of StreamReader
 57        will require keys that (for example) allow it to authenticate with the 3rd party.
 58
 59        Therefore, concrete implementations of AbstractFileBasedStreamReader's config setter should assert that `value` is of the correct
 60        config type for that type of StreamReader.
 61        """
 62        ...
 63
 64    @abstractmethod
 65    def open_file(
 66        self, file: RemoteFile, mode: FileReadMode, encoding: Optional[str], logger: logging.Logger
 67    ) -> IOBase:
 68        """
 69        Return a file handle for reading.
 70
 71        Many sources will be able to use smart_open to implement this method,
 72        for example:
 73
 74        client = boto3.Session(...)
 75        return smart_open.open(remote_file.uri, transport_params={"client": client})
 76        """
 77        ...
 78
 79    @abstractmethod
 80    def get_matching_files(
 81        self,
 82        globs: List[str],
 83        prefix: Optional[str],
 84        logger: logging.Logger,
 85    ) -> Iterable[RemoteFile]:
 86        """
 87        Return all files that match any of the globs.
 88
 89        Example:
 90
 91        The source has files "a.json", "foo/a.json", "foo/bar/a.json"
 92
 93        If globs = ["*.json"] then this method returns ["a.json"].
 94
 95        If globs = ["foo/*.json"] then this method returns ["foo/a.json"].
 96
 97        Utility method `self.filter_files_by_globs` and `self.get_prefixes_from_globs`
 98        are available, which may be helpful when implementing this method.
 99        """
100        ...
101
102    def filter_files_by_globs_and_start_date(
103        self, files: List[RemoteFile], globs: List[str]
104    ) -> Iterable[RemoteFile]:
105        """
106        Utility method for filtering files based on globs.
107        """
108        start_date = (
109            datetime.strptime(self.config.start_date, self.DATE_TIME_FORMAT)
110            if self.config and self.config.start_date
111            else None
112        )
113        seen = set()
114
115        for file in files:
116            if self.file_matches_globs(file, globs):
117                if file.uri not in seen and (not start_date or file.last_modified >= start_date):
118                    seen.add(file.uri)
119                    yield file
120
121    @staticmethod
122    def file_matches_globs(file: RemoteFile, globs: List[str]) -> bool:
123        # Use the GLOBSTAR flag to enable recursive ** matching
124        # (https://facelessuser.github.io/wcmatch/wcmatch/#globstar)
125        return any(globmatch(file.uri, g, flags=GLOBSTAR) for g in globs)
126
127    @staticmethod
128    def get_prefixes_from_globs(globs: List[str]) -> Set[str]:
129        """
130        Utility method for extracting prefixes from the globs.
131        """
132        prefixes = {glob.split("*")[0] for glob in globs}
133        return set(filter(lambda x: bool(x), prefixes))
134
135    def use_file_transfer(self) -> bool:
136        if self.config:
137            return use_file_transfer(self.config)
138        return False
139
140    def preserve_directory_structure(self) -> bool:
141        # fall back to preserve subdirectories if config is not present or incomplete
142        if self.config:
143            return preserve_directory_structure(self.config)
144        return True
145
146    def include_identities_stream(self) -> bool:
147        if self.config:
148            return include_identities_stream(self.config)
149        return False
150
151    def upload(
152        self, file: UploadableRemoteFile, local_directory: str, logger: logging.Logger
153    ) -> Tuple[FileRecordData, AirbyteRecordMessageFileReference]:
154        """
155        This is required for connectors that will support writing to
156        files. It will handle the logic to download,get,read,acquire or
157        whatever is more efficient to get a file from the source.
158
159        Args:
160               file (RemoteFile): The remote file object containing URI and metadata.
161               local_directory (str): The local directory path where the file will be downloaded.
162               logger (logging.Logger): Logger for logging information and errors.
163
164           Returns:
165               AirbyteRecordMessageFileReference: A file reference object containing:
166                   - staging_file_url (str): The absolute path to the referenced file in the staging area.
167                   - file_size_bytes (int): The size of the referenced file in bytes.
168                   - source_file_relative_path (str): The relative path to the referenced file in source.
169        """
170        if not isinstance(file, UploadableRemoteFile):
171            raise TypeError(f"Expected UploadableRemoteFile, got {type(file)}")
172
173        file_size = file.size
174
175        if file_size > self.FILE_SIZE_LIMIT:
176            message = (
177                f"File size exceeds the {self.FILE_SIZE_LIMIT / 1e9} GB limit. File URI: {file.uri}"
178            )
179            raise FileSizeLimitError(
180                message=message, internal_message=message, failure_type=FailureType.config_error
181            )
182
183        file_paths = self._get_file_transfer_paths(
184            source_file_relative_path=file.source_file_relative_path,
185            staging_directory=local_directory,
186        )
187        local_file_path = file_paths[self.LOCAL_FILE_PATH]
188        file_relative_path = file_paths[self.FILE_RELATIVE_PATH]
189        file_name = file_paths[self.FILE_NAME]
190
191        logger.info(
192            f"Starting to download the file {file.file_uri_for_logging} with size: {file_size / (1024 * 1024):,.2f} MB ({file_size / (1024 * 1024 * 1024):.2f} GB)"
193        )
194        start_download_time = time.time()
195
196        file.download_to_local_directory(local_file_path)
197
198        write_duration = time.time() - start_download_time
199        logger.info(
200            f"Finished downloading the file {file.file_uri_for_logging} and saved to {local_file_path} in {write_duration:,.2f} seconds."
201        )
202
203        file_record_data = FileRecordData(
204            folder=file_paths[self.FILE_FOLDER],
205            file_name=file_name,
206            bytes=file_size,
207            id=file.id,
208            mime_type=file.mime_type,
209            created_at=file.created_at,
210            updated_at=file.updated_at,
211            source_uri=file.uri,
212        )
213        file_reference = AirbyteRecordMessageFileReference(
214            staging_file_url=local_file_path,
215            source_file_relative_path=file_relative_path,
216            file_size_bytes=file_size,
217        )
218        return file_record_data, file_reference
219
220    def _get_file_transfer_paths(
221        self, source_file_relative_path: str, staging_directory: str
222    ) -> MutableMapping[str, Any]:
223        """
224        This method is used to get the file transfer paths for a given source file relative path and local directory.
225        It returns a dictionary with the following keys:
226            - FILE_RELATIVE_PATH: The relative path to file in reference to the staging directory.
227            - LOCAL_FILE_PATH: The absolute path to the file.
228            - FILE_NAME: The name of the referenced file.
229            - FILE_FOLDER: The folder of the referenced file.
230        """
231        preserve_directory_structure = self.preserve_directory_structure()
232
233        file_name = path.basename(source_file_relative_path)
234        file_folder = path.dirname(source_file_relative_path)
235        if preserve_directory_structure:
236            # Remove left slashes from source path format to make relative path for writing locally
237            file_relative_path = source_file_relative_path.lstrip("/")
238        else:
239            file_relative_path = file_name
240        local_file_path = path.join(staging_directory, file_relative_path)
241        # Ensure the local directory exists
242        makedirs(path.dirname(local_file_path), exist_ok=True)
243
244        file_paths = {
245            self.FILE_RELATIVE_PATH: file_relative_path,
246            self.LOCAL_FILE_PATH: local_file_path,
247            self.FILE_NAME: file_name,
248            self.FILE_FOLDER: file_folder,
249        }
250        return file_paths

 85class CsvFormat(BaseModel):
 86    class Config(OneOfOptionConfig):
 87        title = "CSV Format"
 88        discriminator = "filetype"
 89
 90    filetype: str = Field(
 91        "csv",
 92        const=True,
 93    )
 94    delimiter: str = Field(
 95        title="Delimiter",
 96        description="The character delimiting individual cells in the CSV data. This may only be a 1-character string. For tab-delimited data enter '\\t'.",
 97        default=",",
 98    )
 99    quote_char: str = Field(
100        title="Quote Character",
101        default='"',
102        description="The character used for quoting CSV values. To disallow quoting, make this field blank.",
103    )
104    escape_char: Optional[str] = Field(
105        title="Escape Character",
106        default=None,
107        description="The character used for escaping special characters. To disallow escaping, leave this field blank.",
108    )
109    encoding: Optional[str] = Field(
110        default="utf8",
111        description='The character encoding of the CSV data. Leave blank to default to <strong>UTF8</strong>. See <a href="https://docs.python.org/3/library/codecs.html#standard-encodings" target="_blank">list of python encodings</a> for allowable options.',
112    )
113    double_quote: bool = Field(
114        title="Double Quote",
115        default=True,
116        description="Whether two quotes in a quoted CSV value denote a single quote in the data.",
117    )
118    null_values: Set[str] = Field(
119        title="Null Values",
120        default=[],
121        description="A set of case-sensitive strings that should be interpreted as null values. For example, if the value 'NA' should be interpreted as null, enter 'NA' in this field.",
122    )
123    strings_can_be_null: bool = Field(
124        title="Strings Can Be Null",
125        default=True,
126        description="Whether strings can be interpreted as null values. If true, strings that match the null_values set will be interpreted as null. If false, strings that match the null_values set will be interpreted as the string itself.",
127    )
128    skip_rows_before_header: int = Field(
129        title="Skip Rows Before Header",
130        default=0,
131        description="The number of rows to skip before the header row. For example, if the header row is on the 3rd row, enter 2 in this field.",
132    )
133    skip_rows_after_header: int = Field(
134        title="Skip Rows After Header",
135        default=0,
136        description="The number of rows to skip after the header row.",
137    )
138    header_definition: Union[CsvHeaderFromCsv, CsvHeaderAutogenerated, CsvHeaderUserProvided] = (
139        Field(
140            title="CSV Header Definition",
141            default=CsvHeaderFromCsv(header_definition_type=CsvHeaderDefinitionType.FROM_CSV.value),
142            description="How headers will be defined. `User Provided` assumes the CSV does not have a header row and uses the headers provided and `Autogenerated` assumes the CSV does not have a header row and the CDK will generate headers using for `f{i}` where `i` is the index starting from 0. Else, the default behavior is to use the header from the CSV file. If a user wants to autogenerate or provide column names for a CSV having headers, they can skip rows.",
143        )
144    )
145    true_values: Set[str] = Field(
146        title="True Values",
147        default=DEFAULT_TRUE_VALUES,
148        description="A set of case-sensitive strings that should be interpreted as true values.",
149    )
150    false_values: Set[str] = Field(
151        title="False Values",
152        default=DEFAULT_FALSE_VALUES,
153        description="A set of case-sensitive strings that should be interpreted as false values.",
154    )
155    inference_type: InferenceType = Field(
156        title="Inference Type",
157        default=InferenceType.NONE,
158        description="How to infer the types of the columns. If none, inference default to strings.",
159        airbyte_hidden=True,
160    )
161    ignore_errors_on_fields_mismatch: bool = Field(
162        title="Ignore errors on field mismatch",
163        default=False,
164        description="Whether to ignore errors that occur when the number of fields in the CSV does not match the number of columns in the schema.",
165    )
166
167    @validator("delimiter")
168    def validate_delimiter(cls, v: str) -> str:
169        if v == r"\t":
170            v = "\t"
171        if len(v) != 1:
172            raise ValueError("delimiter should only be one character")
173        if v in {"\r", "\n"}:
174            raise ValueError(f"delimiter cannot be {v}")
175        return v
176
177    @validator("quote_char")
178    def validate_quote_char(cls, v: str) -> str:
179        if len(v) != 1:
180            raise ValueError("quote_char should only be one character")
181        return v
182
183    @validator("escape_char")
184    def validate_escape_char(cls, v: str) -> str:
185        if v is not None and len(v) != 1:
186            raise ValueError("escape_char should only be one character")
187        return v
188
189    @validator("encoding")
190    def validate_encoding(cls, v: str) -> str:
191        try:
192            codecs.lookup(v)
193        except LookupError:
194            raise ValueError(f"invalid encoding format: {v}")
195        return v
196
197    @root_validator
198    def validate_optional_args(cls, values: Dict[str, Any]) -> Dict[str, Any]:
199        definition_type = values.get("header_definition_type")
200        column_names = values.get("user_provided_column_names")
201        if definition_type == CsvHeaderDefinitionType.USER_PROVIDED and not column_names:
202            raise ValidationError(
203                "`user_provided_column_names` should be defined if the definition 'User Provided'.",
204                model=CsvFormat,
205            )
206        if definition_type != CsvHeaderDefinitionType.USER_PROVIDED and column_names:
207            raise ValidationError(
208                "`user_provided_column_names` should not be defined if the definition is not 'User Provided'.",
209                model=CsvFormat,
210            )
211        return values

86    class Config(OneOfOptionConfig):
87        title = "CSV Format"
88        discriminator = "filetype"

150class CustomFileBasedException(AirbyteTracedException):
151    """
152    A specialized exception for file-based connectors.
153
154    This exception is designed to bypass the default error handling in the file-based CDK, allowing the use of custom error messages.
155    """
156
157    pass

 18class DefaultFileBasedCursor(AbstractFileBasedCursor):
 19    DEFAULT_DAYS_TO_SYNC_IF_HISTORY_IS_FULL = 3
 20    DEFAULT_MAX_HISTORY_SIZE = 10_000
 21    DATE_TIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
 22    CURSOR_FIELD = "_ab_source_file_last_modified"
 23
 24    def __init__(self, stream_config: FileBasedStreamConfig, **_: Any):
 25        super().__init__(stream_config)  # type: ignore [safe-super]
 26        self._file_to_datetime_history: MutableMapping[str, str] = {}
 27        self._time_window_if_history_is_full = timedelta(
 28            days=stream_config.days_to_sync_if_history_is_full
 29            or self.DEFAULT_DAYS_TO_SYNC_IF_HISTORY_IS_FULL
 30        )
 31
 32        if self._time_window_if_history_is_full <= timedelta():
 33            raise ValueError(
 34                f"days_to_sync_if_history_is_full must be a positive timedelta, got {self._time_window_if_history_is_full}"
 35            )
 36
 37        self._start_time = self._compute_start_time()
 38        self._initial_earliest_file_in_history: Optional[RemoteFile] = None
 39
 40    def set_initial_state(self, value: StreamState) -> None:
 41        self._file_to_datetime_history = value.get("history", {})
 42        self._start_time = self._compute_start_time()
 43        self._initial_earliest_file_in_history = self._compute_earliest_file_in_history()
 44
 45    def add_file(self, file: RemoteFile) -> None:
 46        self._file_to_datetime_history[file.uri] = file.last_modified.strftime(
 47            self.DATE_TIME_FORMAT
 48        )
 49        if len(self._file_to_datetime_history) > self.DEFAULT_MAX_HISTORY_SIZE:
 50            # Get the earliest file based on its last modified date and its uri
 51            oldest_file = self._compute_earliest_file_in_history()
 52            if oldest_file:
 53                del self._file_to_datetime_history[oldest_file.uri]
 54            else:
 55                raise Exception(
 56                    "The history is full but there is no files in the history. This should never happen and might be indicative of a bug in the CDK."
 57                )
 58
 59    def get_state(self) -> StreamState:
 60        state = {"history": self._file_to_datetime_history, self.CURSOR_FIELD: self._get_cursor()}
 61        return state
 62
 63    def _get_cursor(self) -> Optional[str]:
 64        """
 65        Returns the cursor value.
 66
 67        Files are synced in order of last-modified with secondary sort on filename, so the cursor value is
 68        a string joining the last-modified timestamp of the last synced file and the name of the file.
 69        """
 70        if self._file_to_datetime_history.items():
 71            filename, timestamp = max(
 72                self._file_to_datetime_history.items(), key=lambda x: (x[1], x[0])
 73            )
 74            return f"{timestamp}_{filename}"
 75        return None
 76
 77    def _is_history_full(self) -> bool:
 78        """
 79        Returns true if the state's history is full, meaning new entries will start to replace old entries.
 80        """
 81        return len(self._file_to_datetime_history) >= self.DEFAULT_MAX_HISTORY_SIZE
 82
 83    def _should_sync_file(self, file: RemoteFile, logger: logging.Logger) -> bool:
 84        if file.uri in self._file_to_datetime_history:
 85            # If the file's uri is in the history, we should sync the file if it has been modified since it was synced
 86            updated_at_from_history = datetime.strptime(
 87                self._file_to_datetime_history[file.uri], self.DATE_TIME_FORMAT
 88            )
 89            if file.last_modified < updated_at_from_history:
 90                logger.warning(
 91                    f"The file {file.uri}'s last modified date is older than the last time it was synced. This is unexpected. Skipping the file."
 92                )
 93            else:
 94                return file.last_modified > updated_at_from_history
 95            return file.last_modified > updated_at_from_history
 96        if self._is_history_full():
 97            if self._initial_earliest_file_in_history is None:
 98                return True
 99            if file.last_modified > self._initial_earliest_file_in_history.last_modified:
100                # If the history is partial and the file's datetime is strictly greater than the earliest file in the history,
101                # we should sync it
102                return True
103            elif file.last_modified == self._initial_earliest_file_in_history.last_modified:
104                # If the history is partial and the file's datetime is equal to the earliest file in the history,
105                # we should sync it if its uri is strictly greater than the earliest file in the history
106                return file.uri > self._initial_earliest_file_in_history.uri
107            else:
108                # Otherwise, only sync the file if it has been modified since the start of the time window
109                return file.last_modified >= self.get_start_time()
110        else:
111            # The file is not in the history and the history is complete. We know we need to sync the file
112            return True
113
114    def get_files_to_sync(
115        self, all_files: Iterable[RemoteFile], logger: logging.Logger
116    ) -> Iterable[RemoteFile]:
117        if self._is_history_full():
118            logger.warning(
119                f"The state history is full. "
120                f"This sync and future syncs won't be able to use the history to filter out duplicate files. "
121                f"It will instead use the time window of {self._time_window_if_history_is_full} to filter out files."
122            )
123        for f in all_files:
124            if self._should_sync_file(f, logger):
125                yield f
126
127    def get_start_time(self) -> datetime:
128        return self._start_time
129
130    def _compute_earliest_file_in_history(self) -> Optional[RemoteFile]:
131        if self._file_to_datetime_history:
132            filename, last_modified = min(
133                self._file_to_datetime_history.items(), key=lambda f: (f[1], f[0])
134            )
135            return RemoteFile(
136                uri=filename, last_modified=datetime.strptime(last_modified, self.DATE_TIME_FORMAT)
137            )
138        else:
139            return None
140
141    def _compute_start_time(self) -> datetime:
142        if not self._file_to_datetime_history:
143            return datetime.min
144        else:
145            earliest = min(self._file_to_datetime_history.values())
146            earliest_dt = datetime.strptime(earliest, self.DATE_TIME_FORMAT)
147            if self._is_history_full():
148                time_window = datetime.now() - self._time_window_if_history_is_full
149                earliest_dt = min(earliest_dt, time_window)
150            return earliest_dt