airbyte-cdk 0.0.0.dev0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- airbyte_cdk/__init__.py +358 -0
- airbyte_cdk/cli/__init__.py +1 -0
- airbyte_cdk/cli/source_declarative_manifest/__init__.py +5 -0
- airbyte_cdk/cli/source_declarative_manifest/_run.py +236 -0
- airbyte_cdk/cli/source_declarative_manifest/spec.json +17 -0
- airbyte_cdk/config_observation.py +104 -0
- airbyte_cdk/connector.py +123 -0
- airbyte_cdk/connector_builder/README.md +53 -0
- airbyte_cdk/connector_builder/__init__.py +3 -0
- airbyte_cdk/connector_builder/connector_builder_handler.py +121 -0
- airbyte_cdk/connector_builder/main.py +107 -0
- airbyte_cdk/connector_builder/models.py +73 -0
- airbyte_cdk/connector_builder/test_reader/__init__.py +7 -0
- airbyte_cdk/connector_builder/test_reader/helpers.py +689 -0
- airbyte_cdk/connector_builder/test_reader/message_grouper.py +173 -0
- airbyte_cdk/connector_builder/test_reader/reader.py +441 -0
- airbyte_cdk/connector_builder/test_reader/types.py +83 -0
- airbyte_cdk/destinations/__init__.py +8 -0
- airbyte_cdk/destinations/destination.py +154 -0
- airbyte_cdk/destinations/vector_db_based/README.md +37 -0
- airbyte_cdk/destinations/vector_db_based/__init__.py +38 -0
- airbyte_cdk/destinations/vector_db_based/config.py +298 -0
- airbyte_cdk/destinations/vector_db_based/document_processor.py +223 -0
- airbyte_cdk/destinations/vector_db_based/embedder.py +303 -0
- airbyte_cdk/destinations/vector_db_based/indexer.py +78 -0
- airbyte_cdk/destinations/vector_db_based/test_utils.py +63 -0
- airbyte_cdk/destinations/vector_db_based/utils.py +35 -0
- airbyte_cdk/destinations/vector_db_based/writer.py +104 -0
- airbyte_cdk/entrypoint.py +414 -0
- airbyte_cdk/exception_handler.py +56 -0
- airbyte_cdk/logger.py +109 -0
- airbyte_cdk/models/__init__.py +72 -0
- airbyte_cdk/models/airbyte_protocol.py +88 -0
- airbyte_cdk/models/airbyte_protocol_serializers.py +44 -0
- airbyte_cdk/models/well_known_types.py +5 -0
- airbyte_cdk/py.typed +0 -0
- airbyte_cdk/sources/__init__.py +26 -0
- airbyte_cdk/sources/abstract_source.py +326 -0
- airbyte_cdk/sources/concurrent_source/__init__.py +8 -0
- airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py +255 -0
- airbyte_cdk/sources/concurrent_source/concurrent_source.py +165 -0
- airbyte_cdk/sources/concurrent_source/concurrent_source_adapter.py +147 -0
- airbyte_cdk/sources/concurrent_source/partition_generation_completed_sentinel.py +24 -0
- airbyte_cdk/sources/concurrent_source/stream_thread_exception.py +25 -0
- airbyte_cdk/sources/concurrent_source/thread_pool_manager.py +115 -0
- airbyte_cdk/sources/config.py +27 -0
- airbyte_cdk/sources/connector_state_manager.py +161 -0
- airbyte_cdk/sources/declarative/__init__.py +3 -0
- airbyte_cdk/sources/declarative/async_job/__init__.py +0 -0
- airbyte_cdk/sources/declarative/async_job/job.py +52 -0
- airbyte_cdk/sources/declarative/async_job/job_orchestrator.py +525 -0
- airbyte_cdk/sources/declarative/async_job/job_tracker.py +79 -0
- airbyte_cdk/sources/declarative/async_job/repository.py +35 -0
- airbyte_cdk/sources/declarative/async_job/status.py +24 -0
- airbyte_cdk/sources/declarative/async_job/timer.py +39 -0
- airbyte_cdk/sources/declarative/auth/__init__.py +8 -0
- airbyte_cdk/sources/declarative/auth/declarative_authenticator.py +42 -0
- airbyte_cdk/sources/declarative/auth/jwt.py +197 -0
- airbyte_cdk/sources/declarative/auth/oauth.py +293 -0
- airbyte_cdk/sources/declarative/auth/selective_authenticator.py +45 -0
- airbyte_cdk/sources/declarative/auth/token.py +267 -0
- airbyte_cdk/sources/declarative/auth/token_provider.py +82 -0
- airbyte_cdk/sources/declarative/checks/__init__.py +24 -0
- airbyte_cdk/sources/declarative/checks/check_dynamic_stream.py +61 -0
- airbyte_cdk/sources/declarative/checks/check_stream.py +56 -0
- airbyte_cdk/sources/declarative/checks/connection_checker.py +35 -0
- airbyte_cdk/sources/declarative/concurrency_level/__init__.py +7 -0
- airbyte_cdk/sources/declarative/concurrency_level/concurrency_level.py +50 -0
- airbyte_cdk/sources/declarative/concurrent_declarative_source.py +526 -0
- airbyte_cdk/sources/declarative/datetime/__init__.py +3 -0
- airbyte_cdk/sources/declarative/datetime/datetime_parser.py +65 -0
- airbyte_cdk/sources/declarative/datetime/min_max_datetime.py +118 -0
- airbyte_cdk/sources/declarative/declarative_component_schema.yaml +3975 -0
- airbyte_cdk/sources/declarative/declarative_source.py +36 -0
- airbyte_cdk/sources/declarative/declarative_stream.py +241 -0
- airbyte_cdk/sources/declarative/decoders/__init__.py +33 -0
- airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py +218 -0
- airbyte_cdk/sources/declarative/decoders/decoder.py +32 -0
- airbyte_cdk/sources/declarative/decoders/decoder_parser.py +30 -0
- airbyte_cdk/sources/declarative/decoders/json_decoder.py +65 -0
- airbyte_cdk/sources/declarative/decoders/noop_decoder.py +21 -0
- airbyte_cdk/sources/declarative/decoders/pagination_decoder_decorator.py +39 -0
- airbyte_cdk/sources/declarative/decoders/xml_decoder.py +98 -0
- airbyte_cdk/sources/declarative/decoders/zipfile_decoder.py +56 -0
- airbyte_cdk/sources/declarative/exceptions.py +9 -0
- airbyte_cdk/sources/declarative/extractors/__init__.py +21 -0
- airbyte_cdk/sources/declarative/extractors/dpath_extractor.py +86 -0
- airbyte_cdk/sources/declarative/extractors/http_selector.py +37 -0
- airbyte_cdk/sources/declarative/extractors/record_extractor.py +27 -0
- airbyte_cdk/sources/declarative/extractors/record_filter.py +91 -0
- airbyte_cdk/sources/declarative/extractors/record_selector.py +170 -0
- airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py +176 -0
- airbyte_cdk/sources/declarative/extractors/type_transformer.py +55 -0
- airbyte_cdk/sources/declarative/incremental/__init__.py +37 -0
- airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py +497 -0
- airbyte_cdk/sources/declarative/incremental/datetime_based_cursor.py +459 -0
- airbyte_cdk/sources/declarative/incremental/declarative_cursor.py +13 -0
- airbyte_cdk/sources/declarative/incremental/global_substream_cursor.py +357 -0
- airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py +380 -0
- airbyte_cdk/sources/declarative/incremental/per_partition_with_global.py +200 -0
- airbyte_cdk/sources/declarative/incremental/resumable_full_refresh_cursor.py +122 -0
- airbyte_cdk/sources/declarative/interpolation/__init__.py +9 -0
- airbyte_cdk/sources/declarative/interpolation/filters.py +139 -0
- airbyte_cdk/sources/declarative/interpolation/interpolated_boolean.py +66 -0
- airbyte_cdk/sources/declarative/interpolation/interpolated_mapping.py +56 -0
- airbyte_cdk/sources/declarative/interpolation/interpolated_nested_mapping.py +52 -0
- airbyte_cdk/sources/declarative/interpolation/interpolated_string.py +79 -0
- airbyte_cdk/sources/declarative/interpolation/interpolation.py +34 -0
- airbyte_cdk/sources/declarative/interpolation/jinja.py +161 -0
- airbyte_cdk/sources/declarative/interpolation/macros.py +191 -0
- airbyte_cdk/sources/declarative/manifest_declarative_source.py +421 -0
- airbyte_cdk/sources/declarative/migrations/__init__.py +0 -0
- airbyte_cdk/sources/declarative/migrations/legacy_to_per_partition_state_migration.py +98 -0
- airbyte_cdk/sources/declarative/migrations/state_migration.py +24 -0
- airbyte_cdk/sources/declarative/models/__init__.py +2 -0
- airbyte_cdk/sources/declarative/models/declarative_component_schema.py +2503 -0
- airbyte_cdk/sources/declarative/parsers/__init__.py +3 -0
- airbyte_cdk/sources/declarative/parsers/custom_code_compiler.py +157 -0
- airbyte_cdk/sources/declarative/parsers/custom_exceptions.py +21 -0
- airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py +172 -0
- airbyte_cdk/sources/declarative/parsers/manifest_reference_resolver.py +213 -0
- airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +3407 -0
- airbyte_cdk/sources/declarative/partition_routers/__init__.py +29 -0
- airbyte_cdk/sources/declarative/partition_routers/async_job_partition_router.py +65 -0
- airbyte_cdk/sources/declarative/partition_routers/cartesian_product_stream_slicer.py +176 -0
- airbyte_cdk/sources/declarative/partition_routers/list_partition_router.py +121 -0
- airbyte_cdk/sources/declarative/partition_routers/partition_router.py +62 -0
- airbyte_cdk/sources/declarative/partition_routers/single_partition_router.py +63 -0
- airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +437 -0
- airbyte_cdk/sources/declarative/requesters/README.md +56 -0
- airbyte_cdk/sources/declarative/requesters/__init__.py +9 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/__init__.py +25 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/__init__.py +23 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/constant_backoff_strategy.py +45 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/exponential_backoff_strategy.py +45 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/header_helper.py +41 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_time_from_header_backoff_strategy.py +70 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_until_time_from_header_backoff_strategy.py +77 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategy.py +17 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/composite_error_handler.py +101 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/default_error_handler.py +147 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/default_http_response_filter.py +40 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/error_handler.py +17 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/http_response_filter.py +179 -0
- airbyte_cdk/sources/declarative/requesters/http_job_repository.py +350 -0
- airbyte_cdk/sources/declarative/requesters/http_requester.py +433 -0
- airbyte_cdk/sources/declarative/requesters/paginators/__init__.py +21 -0
- airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py +327 -0
- airbyte_cdk/sources/declarative/requesters/paginators/no_pagination.py +76 -0
- airbyte_cdk/sources/declarative/requesters/paginators/paginator.py +65 -0
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/__init__.py +25 -0
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/cursor_pagination_strategy.py +98 -0
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py +102 -0
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/page_increment.py +71 -0
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py +48 -0
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/stop_condition.py +66 -0
- airbyte_cdk/sources/declarative/requesters/request_option.py +117 -0
- airbyte_cdk/sources/declarative/requesters/request_options/__init__.py +23 -0
- airbyte_cdk/sources/declarative/requesters/request_options/datetime_based_request_options_provider.py +92 -0
- airbyte_cdk/sources/declarative/requesters/request_options/default_request_options_provider.py +60 -0
- airbyte_cdk/sources/declarative/requesters/request_options/interpolated_nested_request_input_provider.py +59 -0
- airbyte_cdk/sources/declarative/requesters/request_options/interpolated_request_input_provider.py +68 -0
- airbyte_cdk/sources/declarative/requesters/request_options/interpolated_request_options_provider.py +119 -0
- airbyte_cdk/sources/declarative/requesters/request_options/request_options_provider.py +79 -0
- airbyte_cdk/sources/declarative/requesters/request_path.py +15 -0
- airbyte_cdk/sources/declarative/requesters/requester.py +144 -0
- airbyte_cdk/sources/declarative/resolvers/__init__.py +41 -0
- airbyte_cdk/sources/declarative/resolvers/components_resolver.py +55 -0
- airbyte_cdk/sources/declarative/resolvers/config_components_resolver.py +136 -0
- airbyte_cdk/sources/declarative/resolvers/http_components_resolver.py +112 -0
- airbyte_cdk/sources/declarative/retrievers/__init__.py +19 -0
- airbyte_cdk/sources/declarative/retrievers/async_retriever.py +124 -0
- airbyte_cdk/sources/declarative/retrievers/file_uploader.py +89 -0
- airbyte_cdk/sources/declarative/retrievers/retriever.py +54 -0
- airbyte_cdk/sources/declarative/retrievers/simple_retriever.py +702 -0
- airbyte_cdk/sources/declarative/schema/__init__.py +25 -0
- airbyte_cdk/sources/declarative/schema/default_schema_loader.py +47 -0
- airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py +285 -0
- airbyte_cdk/sources/declarative/schema/inline_schema_loader.py +19 -0
- airbyte_cdk/sources/declarative/schema/json_file_schema_loader.py +92 -0
- airbyte_cdk/sources/declarative/schema/schema_loader.py +17 -0
- airbyte_cdk/sources/declarative/spec/__init__.py +7 -0
- airbyte_cdk/sources/declarative/spec/spec.py +48 -0
- airbyte_cdk/sources/declarative/stream_slicers/__init__.py +7 -0
- airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py +93 -0
- airbyte_cdk/sources/declarative/stream_slicers/stream_slicer.py +25 -0
- airbyte_cdk/sources/declarative/transformations/__init__.py +17 -0
- airbyte_cdk/sources/declarative/transformations/add_fields.py +146 -0
- airbyte_cdk/sources/declarative/transformations/dpath_flatten_fields.py +61 -0
- airbyte_cdk/sources/declarative/transformations/flatten_fields.py +52 -0
- airbyte_cdk/sources/declarative/transformations/keys_replace_transformation.py +61 -0
- airbyte_cdk/sources/declarative/transformations/keys_to_lower_transformation.py +22 -0
- airbyte_cdk/sources/declarative/transformations/keys_to_snake_transformation.py +68 -0
- airbyte_cdk/sources/declarative/transformations/remove_fields.py +75 -0
- airbyte_cdk/sources/declarative/transformations/transformation.py +37 -0
- airbyte_cdk/sources/declarative/types.py +25 -0
- airbyte_cdk/sources/declarative/yaml_declarative_source.py +67 -0
- airbyte_cdk/sources/file_based/README.md +152 -0
- airbyte_cdk/sources/file_based/__init__.py +24 -0
- airbyte_cdk/sources/file_based/availability_strategy/__init__.py +11 -0
- airbyte_cdk/sources/file_based/availability_strategy/abstract_file_based_availability_strategy.py +73 -0
- airbyte_cdk/sources/file_based/availability_strategy/default_file_based_availability_strategy.py +149 -0
- airbyte_cdk/sources/file_based/config/__init__.py +0 -0
- airbyte_cdk/sources/file_based/config/abstract_file_based_spec.py +153 -0
- airbyte_cdk/sources/file_based/config/avro_format.py +25 -0
- airbyte_cdk/sources/file_based/config/csv_format.py +210 -0
- airbyte_cdk/sources/file_based/config/excel_format.py +18 -0
- airbyte_cdk/sources/file_based/config/file_based_stream_config.py +99 -0
- airbyte_cdk/sources/file_based/config/jsonl_format.py +18 -0
- airbyte_cdk/sources/file_based/config/parquet_format.py +25 -0
- airbyte_cdk/sources/file_based/config/unstructured_format.py +102 -0
- airbyte_cdk/sources/file_based/config/validate_config_transfer_modes.py +81 -0
- airbyte_cdk/sources/file_based/discovery_policy/__init__.py +8 -0
- airbyte_cdk/sources/file_based/discovery_policy/abstract_discovery_policy.py +21 -0
- airbyte_cdk/sources/file_based/discovery_policy/default_discovery_policy.py +33 -0
- airbyte_cdk/sources/file_based/exceptions.py +159 -0
- airbyte_cdk/sources/file_based/file_based_source.py +466 -0
- airbyte_cdk/sources/file_based/file_based_stream_permissions_reader.py +123 -0
- airbyte_cdk/sources/file_based/file_based_stream_reader.py +209 -0
- airbyte_cdk/sources/file_based/file_record_data.py +22 -0
- airbyte_cdk/sources/file_based/file_types/__init__.py +37 -0
- airbyte_cdk/sources/file_based/file_types/avro_parser.py +233 -0
- airbyte_cdk/sources/file_based/file_types/csv_parser.py +527 -0
- airbyte_cdk/sources/file_based/file_types/excel_parser.py +196 -0
- airbyte_cdk/sources/file_based/file_types/file_transfer.py +30 -0
- airbyte_cdk/sources/file_based/file_types/file_type_parser.py +86 -0
- airbyte_cdk/sources/file_based/file_types/jsonl_parser.py +145 -0
- airbyte_cdk/sources/file_based/file_types/parquet_parser.py +275 -0
- airbyte_cdk/sources/file_based/file_types/unstructured_parser.py +480 -0
- airbyte_cdk/sources/file_based/remote_file.py +18 -0
- airbyte_cdk/sources/file_based/schema_helpers.py +281 -0
- airbyte_cdk/sources/file_based/schema_validation_policies/__init__.py +17 -0
- airbyte_cdk/sources/file_based/schema_validation_policies/abstract_schema_validation_policy.py +20 -0
- airbyte_cdk/sources/file_based/schema_validation_policies/default_schema_validation_policies.py +52 -0
- airbyte_cdk/sources/file_based/stream/__init__.py +13 -0
- airbyte_cdk/sources/file_based/stream/abstract_file_based_stream.py +197 -0
- airbyte_cdk/sources/file_based/stream/concurrent/__init__.py +0 -0
- airbyte_cdk/sources/file_based/stream/concurrent/adapters.py +343 -0
- airbyte_cdk/sources/file_based/stream/concurrent/cursor/__init__.py +9 -0
- airbyte_cdk/sources/file_based/stream/concurrent/cursor/abstract_concurrent_file_based_cursor.py +59 -0
- airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_concurrent_cursor.py +313 -0
- airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_final_state_cursor.py +83 -0
- airbyte_cdk/sources/file_based/stream/cursor/__init__.py +4 -0
- airbyte_cdk/sources/file_based/stream/cursor/abstract_file_based_cursor.py +66 -0
- airbyte_cdk/sources/file_based/stream/cursor/default_file_based_cursor.py +149 -0
- airbyte_cdk/sources/file_based/stream/default_file_based_stream.py +396 -0
- airbyte_cdk/sources/file_based/stream/identities_stream.py +49 -0
- airbyte_cdk/sources/file_based/stream/permissions_file_based_stream.py +92 -0
- airbyte_cdk/sources/file_based/types.py +10 -0
- airbyte_cdk/sources/http_config.py +10 -0
- airbyte_cdk/sources/http_logger.py +55 -0
- airbyte_cdk/sources/message/__init__.py +19 -0
- airbyte_cdk/sources/message/repository.py +137 -0
- airbyte_cdk/sources/source.py +95 -0
- airbyte_cdk/sources/specs/transfer_modes.py +26 -0
- airbyte_cdk/sources/streams/__init__.py +8 -0
- airbyte_cdk/sources/streams/availability_strategy.py +84 -0
- airbyte_cdk/sources/streams/call_rate.py +704 -0
- airbyte_cdk/sources/streams/checkpoint/__init__.py +26 -0
- airbyte_cdk/sources/streams/checkpoint/checkpoint_reader.py +335 -0
- airbyte_cdk/sources/streams/checkpoint/cursor.py +77 -0
- airbyte_cdk/sources/streams/checkpoint/per_partition_key_serializer.py +22 -0
- airbyte_cdk/sources/streams/checkpoint/resumable_full_refresh_cursor.py +51 -0
- airbyte_cdk/sources/streams/checkpoint/substream_resumable_full_refresh_cursor.py +110 -0
- airbyte_cdk/sources/streams/concurrent/README.md +7 -0
- airbyte_cdk/sources/streams/concurrent/__init__.py +3 -0
- airbyte_cdk/sources/streams/concurrent/abstract_stream.py +96 -0
- airbyte_cdk/sources/streams/concurrent/abstract_stream_facade.py +37 -0
- airbyte_cdk/sources/streams/concurrent/adapters.py +397 -0
- airbyte_cdk/sources/streams/concurrent/availability_strategy.py +94 -0
- airbyte_cdk/sources/streams/concurrent/clamping.py +99 -0
- airbyte_cdk/sources/streams/concurrent/cursor.py +481 -0
- airbyte_cdk/sources/streams/concurrent/cursor_types.py +32 -0
- airbyte_cdk/sources/streams/concurrent/default_stream.py +102 -0
- airbyte_cdk/sources/streams/concurrent/exceptions.py +18 -0
- airbyte_cdk/sources/streams/concurrent/helpers.py +42 -0
- airbyte_cdk/sources/streams/concurrent/partition_enqueuer.py +64 -0
- airbyte_cdk/sources/streams/concurrent/partition_reader.py +45 -0
- airbyte_cdk/sources/streams/concurrent/partitions/__init__.py +3 -0
- airbyte_cdk/sources/streams/concurrent/partitions/partition.py +48 -0
- airbyte_cdk/sources/streams/concurrent/partitions/partition_generator.py +18 -0
- airbyte_cdk/sources/streams/concurrent/partitions/stream_slicer.py +21 -0
- airbyte_cdk/sources/streams/concurrent/partitions/types.py +38 -0
- airbyte_cdk/sources/streams/concurrent/state_converters/__init__.py +0 -0
- airbyte_cdk/sources/streams/concurrent/state_converters/abstract_stream_state_converter.py +182 -0
- airbyte_cdk/sources/streams/concurrent/state_converters/datetime_stream_state_converter.py +223 -0
- airbyte_cdk/sources/streams/concurrent/state_converters/incrementing_count_stream_state_converter.py +92 -0
- airbyte_cdk/sources/streams/core.py +703 -0
- airbyte_cdk/sources/streams/http/__init__.py +10 -0
- airbyte_cdk/sources/streams/http/availability_strategy.py +54 -0
- airbyte_cdk/sources/streams/http/error_handlers/__init__.py +22 -0
- airbyte_cdk/sources/streams/http/error_handlers/backoff_strategy.py +28 -0
- airbyte_cdk/sources/streams/http/error_handlers/default_backoff_strategy.py +17 -0
- airbyte_cdk/sources/streams/http/error_handlers/default_error_mapping.py +86 -0
- airbyte_cdk/sources/streams/http/error_handlers/error_handler.py +42 -0
- airbyte_cdk/sources/streams/http/error_handlers/error_message_parser.py +19 -0
- airbyte_cdk/sources/streams/http/error_handlers/http_status_error_handler.py +110 -0
- airbyte_cdk/sources/streams/http/error_handlers/json_error_message_parser.py +52 -0
- airbyte_cdk/sources/streams/http/error_handlers/response_models.py +65 -0
- airbyte_cdk/sources/streams/http/exceptions.py +61 -0
- airbyte_cdk/sources/streams/http/http.py +673 -0
- airbyte_cdk/sources/streams/http/http_client.py +531 -0
- airbyte_cdk/sources/streams/http/rate_limiting.py +158 -0
- airbyte_cdk/sources/streams/http/requests_native_auth/__init__.py +14 -0
- airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py +479 -0
- airbyte_cdk/sources/streams/http/requests_native_auth/abstract_token.py +34 -0
- airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py +436 -0
- airbyte_cdk/sources/streams/http/requests_native_auth/token.py +83 -0
- airbyte_cdk/sources/streams/permissions/identities_stream.py +75 -0
- airbyte_cdk/sources/streams/utils/__init__.py +3 -0
- airbyte_cdk/sources/types.py +169 -0
- airbyte_cdk/sources/utils/__init__.py +7 -0
- airbyte_cdk/sources/utils/casing.py +12 -0
- airbyte_cdk/sources/utils/files_directory.py +15 -0
- airbyte_cdk/sources/utils/record_helper.py +53 -0
- airbyte_cdk/sources/utils/schema_helpers.py +230 -0
- airbyte_cdk/sources/utils/slice_logger.py +57 -0
- airbyte_cdk/sources/utils/transform.py +277 -0
- airbyte_cdk/sources/utils/types.py +7 -0
- airbyte_cdk/sql/__init__.py +0 -0
- airbyte_cdk/sql/_util/__init__.py +0 -0
- airbyte_cdk/sql/_util/hashing.py +34 -0
- airbyte_cdk/sql/_util/name_normalizers.py +92 -0
- airbyte_cdk/sql/constants.py +32 -0
- airbyte_cdk/sql/exceptions.py +235 -0
- airbyte_cdk/sql/secrets.py +123 -0
- airbyte_cdk/sql/shared/__init__.py +15 -0
- airbyte_cdk/sql/shared/catalog_providers.py +145 -0
- airbyte_cdk/sql/shared/sql_processor.py +786 -0
- airbyte_cdk/sql/types.py +160 -0
- airbyte_cdk/test/__init__.py +7 -0
- airbyte_cdk/test/catalog_builder.py +81 -0
- airbyte_cdk/test/entrypoint_wrapper.py +250 -0
- airbyte_cdk/test/mock_http/__init__.py +6 -0
- airbyte_cdk/test/mock_http/matcher.py +41 -0
- airbyte_cdk/test/mock_http/mocker.py +185 -0
- airbyte_cdk/test/mock_http/request.py +103 -0
- airbyte_cdk/test/mock_http/response.py +28 -0
- airbyte_cdk/test/mock_http/response_builder.py +237 -0
- airbyte_cdk/test/state_builder.py +33 -0
- airbyte_cdk/test/utils/__init__.py +1 -0
- airbyte_cdk/test/utils/data.py +24 -0
- airbyte_cdk/test/utils/http_mocking.py +16 -0
- airbyte_cdk/test/utils/manifest_only_fixtures.py +59 -0
- airbyte_cdk/test/utils/reading.py +26 -0
- airbyte_cdk/utils/__init__.py +10 -0
- airbyte_cdk/utils/airbyte_secrets_utils.py +80 -0
- airbyte_cdk/utils/analytics_message.py +25 -0
- airbyte_cdk/utils/constants.py +5 -0
- airbyte_cdk/utils/datetime_format_inferrer.py +94 -0
- airbyte_cdk/utils/datetime_helpers.py +499 -0
- airbyte_cdk/utils/event_timing.py +85 -0
- airbyte_cdk/utils/is_cloud_environment.py +18 -0
- airbyte_cdk/utils/mapping_helpers.py +162 -0
- airbyte_cdk/utils/message_utils.py +26 -0
- airbyte_cdk/utils/oneof_option_config.py +33 -0
- airbyte_cdk/utils/print_buffer.py +75 -0
- airbyte_cdk/utils/schema_inferrer.py +270 -0
- airbyte_cdk/utils/slice_hasher.py +37 -0
- airbyte_cdk/utils/spec_schema_transformations.py +26 -0
- airbyte_cdk/utils/stream_status_utils.py +43 -0
- airbyte_cdk/utils/traced_exception.py +145 -0
- airbyte_cdk-0.0.0.dev0.dist-info/LICENSE.txt +19 -0
- airbyte_cdk-0.0.0.dev0.dist-info/LICENSE_SHORT +1 -0
- airbyte_cdk-0.0.0.dev0.dist-info/METADATA +111 -0
- airbyte_cdk-0.0.0.dev0.dist-info/RECORD +368 -0
- airbyte_cdk-0.0.0.dev0.dist-info/WHEEL +4 -0
- airbyte_cdk-0.0.0.dev0.dist-info/entry_points.txt +3 -0
|
@@ -0,0 +1,525 @@
|
|
|
1
|
+
# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
|
|
2
|
+
|
|
3
|
+
import logging
|
|
4
|
+
import threading
|
|
5
|
+
import time
|
|
6
|
+
import traceback
|
|
7
|
+
import uuid
|
|
8
|
+
from datetime import timedelta
|
|
9
|
+
from typing import (
|
|
10
|
+
Any,
|
|
11
|
+
Generator,
|
|
12
|
+
Generic,
|
|
13
|
+
Iterable,
|
|
14
|
+
List,
|
|
15
|
+
Mapping,
|
|
16
|
+
Optional,
|
|
17
|
+
Set,
|
|
18
|
+
Tuple,
|
|
19
|
+
Type,
|
|
20
|
+
TypeVar,
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
from airbyte_cdk.logger import lazy_log
|
|
24
|
+
from airbyte_cdk.models import FailureType
|
|
25
|
+
from airbyte_cdk.sources.declarative.async_job.job import AsyncJob
|
|
26
|
+
from airbyte_cdk.sources.declarative.async_job.job_tracker import (
|
|
27
|
+
ConcurrentJobLimitReached,
|
|
28
|
+
JobTracker,
|
|
29
|
+
)
|
|
30
|
+
from airbyte_cdk.sources.declarative.async_job.repository import AsyncJobRepository
|
|
31
|
+
from airbyte_cdk.sources.declarative.async_job.status import AsyncJobStatus
|
|
32
|
+
from airbyte_cdk.sources.message import MessageRepository
|
|
33
|
+
from airbyte_cdk.sources.types import StreamSlice
|
|
34
|
+
from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
|
|
35
|
+
from airbyte_cdk.utils.traced_exception import AirbyteTracedException
|
|
36
|
+
|
|
37
|
+
LOGGER = logging.getLogger("airbyte")
|
|
38
|
+
_NO_TIMEOUT = timedelta.max
|
|
39
|
+
_API_SIDE_RUNNING_STATUS = {AsyncJobStatus.RUNNING, AsyncJobStatus.TIMED_OUT}
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
class AsyncPartition:
|
|
43
|
+
"""
|
|
44
|
+
This bucket of api_jobs is a bit useless for this iteration but should become interesting when we will be able to split jobs
|
|
45
|
+
"""
|
|
46
|
+
|
|
47
|
+
_MAX_NUMBER_OF_ATTEMPTS = 3
|
|
48
|
+
|
|
49
|
+
def __init__(self, jobs: List[AsyncJob], stream_slice: StreamSlice) -> None:
|
|
50
|
+
self._attempts_per_job = {job: 1 for job in jobs}
|
|
51
|
+
self._stream_slice = stream_slice
|
|
52
|
+
|
|
53
|
+
def has_reached_max_attempt(self) -> bool:
|
|
54
|
+
return any(
|
|
55
|
+
map(
|
|
56
|
+
lambda attempt_count: attempt_count >= self._MAX_NUMBER_OF_ATTEMPTS,
|
|
57
|
+
self._attempts_per_job.values(),
|
|
58
|
+
)
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
def replace_job(self, job_to_replace: AsyncJob, new_jobs: List[AsyncJob]) -> None:
|
|
62
|
+
current_attempt_count = self._attempts_per_job.pop(job_to_replace, None)
|
|
63
|
+
if current_attempt_count is None:
|
|
64
|
+
raise ValueError("Could not find job to replace")
|
|
65
|
+
elif current_attempt_count >= self._MAX_NUMBER_OF_ATTEMPTS:
|
|
66
|
+
raise ValueError(f"Max attempt reached for job in partition {self._stream_slice}")
|
|
67
|
+
|
|
68
|
+
new_attempt_count = current_attempt_count + 1
|
|
69
|
+
for job in new_jobs:
|
|
70
|
+
self._attempts_per_job[job] = new_attempt_count
|
|
71
|
+
|
|
72
|
+
def should_split(self, job: AsyncJob) -> bool:
|
|
73
|
+
"""
|
|
74
|
+
Not used right now but once we support job split, we should split based on the number of attempts
|
|
75
|
+
"""
|
|
76
|
+
return False
|
|
77
|
+
|
|
78
|
+
@property
|
|
79
|
+
def jobs(self) -> Iterable[AsyncJob]:
|
|
80
|
+
return self._attempts_per_job.keys()
|
|
81
|
+
|
|
82
|
+
@property
|
|
83
|
+
def stream_slice(self) -> StreamSlice:
|
|
84
|
+
return self._stream_slice
|
|
85
|
+
|
|
86
|
+
@property
|
|
87
|
+
def status(self) -> AsyncJobStatus:
|
|
88
|
+
"""
|
|
89
|
+
Given different job statuses, the priority is: FAILED, TIMED_OUT, RUNNING. Else, it means everything is completed.
|
|
90
|
+
"""
|
|
91
|
+
statuses = set(map(lambda job: job.status(), self.jobs))
|
|
92
|
+
if statuses == {AsyncJobStatus.COMPLETED}:
|
|
93
|
+
return AsyncJobStatus.COMPLETED
|
|
94
|
+
elif AsyncJobStatus.FAILED in statuses:
|
|
95
|
+
return AsyncJobStatus.FAILED
|
|
96
|
+
elif AsyncJobStatus.TIMED_OUT in statuses:
|
|
97
|
+
return AsyncJobStatus.TIMED_OUT
|
|
98
|
+
else:
|
|
99
|
+
return AsyncJobStatus.RUNNING
|
|
100
|
+
|
|
101
|
+
def __repr__(self) -> str:
|
|
102
|
+
return f"AsyncPartition(stream_slice={self._stream_slice}, attempt_per_job={self._attempts_per_job})"
|
|
103
|
+
|
|
104
|
+
def __json_serializable__(self) -> Any:
|
|
105
|
+
return self._stream_slice
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
T = TypeVar("T")
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
class LookaheadIterator(Generic[T]):
|
|
112
|
+
def __init__(self, iterable: Iterable[T]) -> None:
|
|
113
|
+
self._iterator = iter(iterable)
|
|
114
|
+
self._buffer: List[T] = []
|
|
115
|
+
|
|
116
|
+
def __iter__(self) -> "LookaheadIterator[T]":
|
|
117
|
+
return self
|
|
118
|
+
|
|
119
|
+
def __next__(self) -> T:
|
|
120
|
+
if self._buffer:
|
|
121
|
+
return self._buffer.pop()
|
|
122
|
+
else:
|
|
123
|
+
return next(self._iterator)
|
|
124
|
+
|
|
125
|
+
def has_next(self) -> bool:
|
|
126
|
+
if self._buffer:
|
|
127
|
+
return True
|
|
128
|
+
|
|
129
|
+
try:
|
|
130
|
+
self._buffer = [next(self._iterator)]
|
|
131
|
+
except StopIteration:
|
|
132
|
+
return False
|
|
133
|
+
else:
|
|
134
|
+
return True
|
|
135
|
+
|
|
136
|
+
def add_at_the_beginning(self, item: T) -> None:
|
|
137
|
+
self._buffer = [item] + self._buffer
|
|
138
|
+
|
|
139
|
+
|
|
140
|
+
class AsyncJobOrchestrator:
|
|
141
|
+
_WAIT_TIME_BETWEEN_STATUS_UPDATE_IN_SECONDS = 5
|
|
142
|
+
_KNOWN_JOB_STATUSES = {
|
|
143
|
+
AsyncJobStatus.COMPLETED,
|
|
144
|
+
AsyncJobStatus.FAILED,
|
|
145
|
+
AsyncJobStatus.RUNNING,
|
|
146
|
+
AsyncJobStatus.TIMED_OUT,
|
|
147
|
+
}
|
|
148
|
+
_RUNNING_ON_API_SIDE_STATUS = {AsyncJobStatus.RUNNING, AsyncJobStatus.TIMED_OUT}
|
|
149
|
+
|
|
150
|
+
def __init__(
|
|
151
|
+
self,
|
|
152
|
+
job_repository: AsyncJobRepository,
|
|
153
|
+
slices: Iterable[StreamSlice],
|
|
154
|
+
job_tracker: JobTracker,
|
|
155
|
+
message_repository: MessageRepository,
|
|
156
|
+
exceptions_to_break_on: Iterable[Type[Exception]] = tuple(),
|
|
157
|
+
has_bulk_parent: bool = False,
|
|
158
|
+
) -> None:
|
|
159
|
+
"""
|
|
160
|
+
If the stream slices provided as a parameters relies on a async job streams that relies on the same JobTracker, `has_bulk_parent`
|
|
161
|
+
needs to be set to True as jobs creation needs to be prioritized on the parent level. Doing otherwise could lead to a situation
|
|
162
|
+
where the child has taken up all the job budget without room to the parent to create more which would lead to an infinite loop of
|
|
163
|
+
"trying to start a parent job" and "ConcurrentJobLimitReached".
|
|
164
|
+
"""
|
|
165
|
+
if {*AsyncJobStatus} != self._KNOWN_JOB_STATUSES:
|
|
166
|
+
# this is to prevent developers updating the possible statuses without updating the logic of this class
|
|
167
|
+
raise ValueError(
|
|
168
|
+
"An AsyncJobStatus has been either removed or added which means the logic of this class needs to be reviewed. Once the logic has been updated, please update _KNOWN_JOB_STATUSES"
|
|
169
|
+
)
|
|
170
|
+
|
|
171
|
+
self._job_repository: AsyncJobRepository = job_repository
|
|
172
|
+
self._slice_iterator = LookaheadIterator(slices)
|
|
173
|
+
self._running_partitions: List[AsyncPartition] = []
|
|
174
|
+
self._job_tracker = job_tracker
|
|
175
|
+
self._message_repository = message_repository
|
|
176
|
+
self._exceptions_to_break_on: Tuple[Type[Exception], ...] = tuple(exceptions_to_break_on)
|
|
177
|
+
self._has_bulk_parent = has_bulk_parent
|
|
178
|
+
|
|
179
|
+
self._non_breaking_exceptions: List[Exception] = []
|
|
180
|
+
|
|
181
|
+
def _replace_failed_jobs(self, partition: AsyncPartition) -> None:
|
|
182
|
+
failed_status_jobs = (AsyncJobStatus.FAILED,)
|
|
183
|
+
jobs_to_replace = [job for job in partition.jobs if job.status() in failed_status_jobs]
|
|
184
|
+
for job in jobs_to_replace:
|
|
185
|
+
new_job = self._start_job(job.job_parameters(), job.api_job_id())
|
|
186
|
+
partition.replace_job(job, [new_job])
|
|
187
|
+
|
|
188
|
+
def _start_jobs(self) -> None:
|
|
189
|
+
"""
|
|
190
|
+
Retry failed jobs and start jobs for each slice in the slice iterator.
|
|
191
|
+
This method iterates over the running jobs and slice iterator and starts a job for each slice.
|
|
192
|
+
The started jobs are added to the running partitions.
|
|
193
|
+
Returns:
|
|
194
|
+
None
|
|
195
|
+
|
|
196
|
+
However, the first iteration is for sendgrid which only has one job.
|
|
197
|
+
"""
|
|
198
|
+
at_least_one_slice_consumed_from_slice_iterator_during_current_iteration = False
|
|
199
|
+
_slice = None
|
|
200
|
+
try:
|
|
201
|
+
for partition in self._running_partitions:
|
|
202
|
+
self._replace_failed_jobs(partition)
|
|
203
|
+
|
|
204
|
+
if (
|
|
205
|
+
self._has_bulk_parent
|
|
206
|
+
and self._running_partitions
|
|
207
|
+
and self._slice_iterator.has_next()
|
|
208
|
+
):
|
|
209
|
+
LOGGER.debug(
|
|
210
|
+
"This AsyncJobOrchestrator is operating as a child of a bulk stream hence we limit the number of concurrent jobs on the child until there are no more parent slices to avoid the child taking all the API job budget"
|
|
211
|
+
)
|
|
212
|
+
return
|
|
213
|
+
|
|
214
|
+
for _slice in self._slice_iterator:
|
|
215
|
+
at_least_one_slice_consumed_from_slice_iterator_during_current_iteration = True
|
|
216
|
+
job = self._start_job(_slice)
|
|
217
|
+
self._running_partitions.append(AsyncPartition([job], _slice))
|
|
218
|
+
if self._has_bulk_parent and self._slice_iterator.has_next():
|
|
219
|
+
break
|
|
220
|
+
except ConcurrentJobLimitReached:
|
|
221
|
+
if at_least_one_slice_consumed_from_slice_iterator_during_current_iteration:
|
|
222
|
+
# this means a slice has been consumed but the job couldn't be create therefore we need to put it back at the beginning of the _slice_iterator
|
|
223
|
+
self._slice_iterator.add_at_the_beginning(_slice) # type: ignore # we know it's not None here because `ConcurrentJobLimitReached` happens during the for loop
|
|
224
|
+
LOGGER.debug(
|
|
225
|
+
"Waiting before creating more jobs as the limit of concurrent jobs has been reached. Will try again later..."
|
|
226
|
+
)
|
|
227
|
+
|
|
228
|
+
def _start_job(self, _slice: StreamSlice, previous_job_id: Optional[str] = None) -> AsyncJob:
|
|
229
|
+
if previous_job_id:
|
|
230
|
+
id_to_replace = previous_job_id
|
|
231
|
+
lazy_log(LOGGER, logging.DEBUG, lambda: f"Attempting to replace job {id_to_replace}...")
|
|
232
|
+
else:
|
|
233
|
+
id_to_replace = self._job_tracker.try_to_get_intent()
|
|
234
|
+
|
|
235
|
+
try:
|
|
236
|
+
job = self._job_repository.start(_slice)
|
|
237
|
+
self._job_tracker.add_job(id_to_replace, job.api_job_id())
|
|
238
|
+
return job
|
|
239
|
+
except Exception as exception:
|
|
240
|
+
LOGGER.warning(f"Exception has occurred during job creation: {exception}")
|
|
241
|
+
if self._is_breaking_exception(exception):
|
|
242
|
+
self._job_tracker.remove_job(id_to_replace)
|
|
243
|
+
raise exception
|
|
244
|
+
return self._keep_api_budget_with_failed_job(_slice, exception, id_to_replace)
|
|
245
|
+
|
|
246
|
+
def _keep_api_budget_with_failed_job(
|
|
247
|
+
self, _slice: StreamSlice, exception: Exception, intent: str
|
|
248
|
+
) -> AsyncJob:
|
|
249
|
+
"""
|
|
250
|
+
We have a mechanism to retry job. It is used when a job status is FAILED or TIMED_OUT. The easiest way to retry is to have this job
|
|
251
|
+
as created in a failed state and leverage the retry for failed/timed out jobs. This way, we don't have to have another process for
|
|
252
|
+
retrying jobs that couldn't be started.
|
|
253
|
+
"""
|
|
254
|
+
LOGGER.warning(
|
|
255
|
+
f"Could not start job for slice {_slice}. Job will be flagged as failed and retried if max number of attempts not reached: {exception}"
|
|
256
|
+
)
|
|
257
|
+
traced_exception = (
|
|
258
|
+
exception
|
|
259
|
+
if isinstance(exception, AirbyteTracedException)
|
|
260
|
+
else AirbyteTracedException.from_exception(exception)
|
|
261
|
+
)
|
|
262
|
+
# Even though we're not sure this will break the stream, we will emit here for simplicity's sake. If we wanted to be more accurate,
|
|
263
|
+
# we would keep the exceptions in-memory until we know that we have reached the max attempt.
|
|
264
|
+
self._message_repository.emit_message(traced_exception.as_airbyte_message())
|
|
265
|
+
job = self._create_failed_job(_slice)
|
|
266
|
+
self._job_tracker.add_job(intent, job.api_job_id())
|
|
267
|
+
return job
|
|
268
|
+
|
|
269
|
+
def _create_failed_job(self, stream_slice: StreamSlice) -> AsyncJob:
|
|
270
|
+
job = AsyncJob(f"{uuid.uuid4()} - Job that could not start", stream_slice, _NO_TIMEOUT)
|
|
271
|
+
job.update_status(AsyncJobStatus.FAILED)
|
|
272
|
+
return job
|
|
273
|
+
|
|
274
|
+
def _get_running_jobs(self) -> Set[AsyncJob]:
|
|
275
|
+
"""
|
|
276
|
+
Returns a set of running AsyncJob objects.
|
|
277
|
+
|
|
278
|
+
Returns:
|
|
279
|
+
Set[AsyncJob]: A set of AsyncJob objects that are currently running.
|
|
280
|
+
"""
|
|
281
|
+
return {
|
|
282
|
+
job
|
|
283
|
+
for partition in self._running_partitions
|
|
284
|
+
for job in partition.jobs
|
|
285
|
+
if job.status() == AsyncJobStatus.RUNNING
|
|
286
|
+
}
|
|
287
|
+
|
|
288
|
+
def _update_jobs_status(self) -> None:
|
|
289
|
+
"""
|
|
290
|
+
Update the status of all running jobs in the repository.
|
|
291
|
+
"""
|
|
292
|
+
running_jobs = self._get_running_jobs()
|
|
293
|
+
if running_jobs:
|
|
294
|
+
# update the status only if there are RUNNING jobs
|
|
295
|
+
self._job_repository.update_jobs_status(running_jobs)
|
|
296
|
+
|
|
297
|
+
def _wait_on_status_update(self) -> None:
|
|
298
|
+
"""
|
|
299
|
+
Waits for a specified amount of time between status updates.
|
|
300
|
+
|
|
301
|
+
|
|
302
|
+
This method is used to introduce a delay between status updates in order to avoid excessive polling.
|
|
303
|
+
The duration of the delay is determined by the value of `_WAIT_TIME_BETWEEN_STATUS_UPDATE_IN_SECONDS`.
|
|
304
|
+
|
|
305
|
+
Returns:
|
|
306
|
+
None
|
|
307
|
+
"""
|
|
308
|
+
lazy_log(
|
|
309
|
+
LOGGER,
|
|
310
|
+
logging.DEBUG,
|
|
311
|
+
lambda: f"Polling status in progress. There are currently {len(self._running_partitions)} running partitions.",
|
|
312
|
+
)
|
|
313
|
+
|
|
314
|
+
lazy_log(
|
|
315
|
+
LOGGER,
|
|
316
|
+
logging.DEBUG,
|
|
317
|
+
lambda: f"Waiting for {self._WAIT_TIME_BETWEEN_STATUS_UPDATE_IN_SECONDS} seconds before next poll...",
|
|
318
|
+
)
|
|
319
|
+
time.sleep(self._WAIT_TIME_BETWEEN_STATUS_UPDATE_IN_SECONDS)
|
|
320
|
+
|
|
321
|
+
def _process_completed_partition(self, partition: AsyncPartition) -> None:
|
|
322
|
+
"""
|
|
323
|
+
Process a completed partition.
|
|
324
|
+
Args:
|
|
325
|
+
partition (AsyncPartition): The completed partition to process.
|
|
326
|
+
"""
|
|
327
|
+
job_ids = list(map(lambda job: job.api_job_id(), {job for job in partition.jobs}))
|
|
328
|
+
LOGGER.info(
|
|
329
|
+
f"The following jobs for stream slice {partition.stream_slice} have been completed: {job_ids}."
|
|
330
|
+
)
|
|
331
|
+
|
|
332
|
+
# It is important to remove the jobs from the job tracker before yielding the partition as the caller might try to schedule jobs
|
|
333
|
+
# but won't be able to as all jobs slots are taken even though job is done.
|
|
334
|
+
for job in partition.jobs:
|
|
335
|
+
self._job_tracker.remove_job(job.api_job_id())
|
|
336
|
+
|
|
337
|
+
def _process_running_partitions_and_yield_completed_ones(
|
|
338
|
+
self,
|
|
339
|
+
) -> Generator[AsyncPartition, Any, None]:
|
|
340
|
+
"""
|
|
341
|
+
Process the running partitions.
|
|
342
|
+
|
|
343
|
+
Yields:
|
|
344
|
+
AsyncPartition: The processed partition.
|
|
345
|
+
|
|
346
|
+
Raises:
|
|
347
|
+
Any: Any exception raised during processing.
|
|
348
|
+
"""
|
|
349
|
+
current_running_partitions: List[AsyncPartition] = []
|
|
350
|
+
for partition in self._running_partitions:
|
|
351
|
+
match partition.status:
|
|
352
|
+
case AsyncJobStatus.COMPLETED:
|
|
353
|
+
self._process_completed_partition(partition)
|
|
354
|
+
yield partition
|
|
355
|
+
case AsyncJobStatus.RUNNING:
|
|
356
|
+
current_running_partitions.append(partition)
|
|
357
|
+
case _ if partition.has_reached_max_attempt():
|
|
358
|
+
self._stop_partition(partition)
|
|
359
|
+
self._process_partitions_with_errors(partition)
|
|
360
|
+
case _:
|
|
361
|
+
self._stop_timed_out_jobs(partition)
|
|
362
|
+
# re-allocate FAILED jobs, but TIMEOUT jobs are not re-allocated
|
|
363
|
+
self._reallocate_partition(current_running_partitions, partition)
|
|
364
|
+
|
|
365
|
+
# We only remove completed / timeout jobs jobs as we want failed jobs to be re-allocated in priority
|
|
366
|
+
self._remove_completed_or_timed_out_jobs(partition)
|
|
367
|
+
|
|
368
|
+
# update the referenced list with running partitions
|
|
369
|
+
self._running_partitions = current_running_partitions
|
|
370
|
+
|
|
371
|
+
def _stop_partition(self, partition: AsyncPartition) -> None:
|
|
372
|
+
for job in partition.jobs:
|
|
373
|
+
if job.status() in _API_SIDE_RUNNING_STATUS:
|
|
374
|
+
self._abort_job(job, free_job_allocation=True)
|
|
375
|
+
else:
|
|
376
|
+
self._job_tracker.remove_job(job.api_job_id())
|
|
377
|
+
|
|
378
|
+
def _stop_timed_out_jobs(self, partition: AsyncPartition) -> None:
|
|
379
|
+
for job in partition.jobs:
|
|
380
|
+
if job.status() == AsyncJobStatus.TIMED_OUT:
|
|
381
|
+
self._abort_job(job, free_job_allocation=True)
|
|
382
|
+
raise AirbyteTracedException(
|
|
383
|
+
internal_message=f"Job {job.api_job_id()} has timed out. Try increasing the `polling job timeout`.",
|
|
384
|
+
failure_type=FailureType.config_error,
|
|
385
|
+
)
|
|
386
|
+
|
|
387
|
+
def _abort_job(self, job: AsyncJob, free_job_allocation: bool = True) -> None:
|
|
388
|
+
try:
|
|
389
|
+
self._job_repository.abort(job)
|
|
390
|
+
if free_job_allocation:
|
|
391
|
+
self._job_tracker.remove_job(job.api_job_id())
|
|
392
|
+
except Exception as exception:
|
|
393
|
+
LOGGER.warning(f"Could not free budget for job {job.api_job_id()}: {exception}")
|
|
394
|
+
|
|
395
|
+
def _remove_completed_or_timed_out_jobs(self, partition: AsyncPartition) -> None:
|
|
396
|
+
"""
|
|
397
|
+
Remove completed or timed out jobs from the partition.
|
|
398
|
+
|
|
399
|
+
Args:
|
|
400
|
+
partition (AsyncPartition): The partition to process.
|
|
401
|
+
"""
|
|
402
|
+
for job in partition.jobs:
|
|
403
|
+
if job.status() in [AsyncJobStatus.COMPLETED, AsyncJobStatus.TIMED_OUT]:
|
|
404
|
+
self._job_tracker.remove_job(job.api_job_id())
|
|
405
|
+
|
|
406
|
+
def _reallocate_partition(
|
|
407
|
+
self,
|
|
408
|
+
current_running_partitions: List[AsyncPartition],
|
|
409
|
+
partition: AsyncPartition,
|
|
410
|
+
) -> None:
|
|
411
|
+
"""
|
|
412
|
+
Reallocate the partition by starting a new job for each job in the
|
|
413
|
+
partition.
|
|
414
|
+
Args:
|
|
415
|
+
current_running_partitions (list): The list of currently running partitions.
|
|
416
|
+
partition (AsyncPartition): The partition to reallocate.
|
|
417
|
+
"""
|
|
418
|
+
for job in partition.jobs:
|
|
419
|
+
if job.status() != AsyncJobStatus.TIMED_OUT:
|
|
420
|
+
# allow the FAILED jobs to be re-allocated for partition
|
|
421
|
+
current_running_partitions.insert(0, partition)
|
|
422
|
+
|
|
423
|
+
def _process_partitions_with_errors(self, partition: AsyncPartition) -> None:
|
|
424
|
+
"""
|
|
425
|
+
Process a partition with status errors (FAILED and TIMEOUT).
|
|
426
|
+
|
|
427
|
+
Args:
|
|
428
|
+
partition (AsyncPartition): The partition to process.
|
|
429
|
+
Returns:
|
|
430
|
+
AirbyteTracedException: An exception indicating that at least one job could not be completed.
|
|
431
|
+
Raises:
|
|
432
|
+
AirbyteTracedException: If at least one job could not be completed.
|
|
433
|
+
"""
|
|
434
|
+
status_by_job_id = {job.api_job_id(): job.status() for job in partition.jobs}
|
|
435
|
+
self._non_breaking_exceptions.append(
|
|
436
|
+
AirbyteTracedException(
|
|
437
|
+
internal_message=f"At least one job could not be completed for slice {partition.stream_slice}. Job statuses were: {status_by_job_id}. See warning logs for more information.",
|
|
438
|
+
failure_type=FailureType.config_error,
|
|
439
|
+
)
|
|
440
|
+
)
|
|
441
|
+
|
|
442
|
+
def create_and_get_completed_partitions(self) -> Iterable[AsyncPartition]:
|
|
443
|
+
"""
|
|
444
|
+
Creates and retrieves completed partitions.
|
|
445
|
+
This method continuously starts jobs, updates job status, processes running partitions,
|
|
446
|
+
logs polling partitions, and waits for status updates. It yields completed partitions
|
|
447
|
+
as they become available.
|
|
448
|
+
|
|
449
|
+
Returns:
|
|
450
|
+
An iterable of completed partitions, represented as AsyncPartition objects.
|
|
451
|
+
Each partition is wrapped in an Optional, allowing for None values.
|
|
452
|
+
"""
|
|
453
|
+
while True:
|
|
454
|
+
try:
|
|
455
|
+
lazy_log(
|
|
456
|
+
LOGGER,
|
|
457
|
+
logging.DEBUG,
|
|
458
|
+
lambda: f"JobOrchestrator loop - (Thread {threading.get_native_id()}, AsyncJobOrchestrator {self}) is starting the async job loop",
|
|
459
|
+
)
|
|
460
|
+
self._start_jobs()
|
|
461
|
+
if not self._slice_iterator.has_next() and not self._running_partitions:
|
|
462
|
+
break
|
|
463
|
+
|
|
464
|
+
self._update_jobs_status()
|
|
465
|
+
yield from self._process_running_partitions_and_yield_completed_ones()
|
|
466
|
+
self._wait_on_status_update()
|
|
467
|
+
except Exception as exception:
|
|
468
|
+
LOGGER.warning(
|
|
469
|
+
f"Caught exception that stops the processing of the jobs: {exception}. Traceback: {traceback.format_exc()}"
|
|
470
|
+
)
|
|
471
|
+
if self._is_breaking_exception(exception):
|
|
472
|
+
self._abort_all_running_jobs()
|
|
473
|
+
raise exception
|
|
474
|
+
|
|
475
|
+
self._non_breaking_exceptions.append(exception)
|
|
476
|
+
|
|
477
|
+
LOGGER.info(
|
|
478
|
+
f"JobOrchestrator loop - Thread (Thread {threading.get_native_id()}, AsyncJobOrchestrator {self}) completed! Errors during creation were {self._non_breaking_exceptions}"
|
|
479
|
+
)
|
|
480
|
+
if self._non_breaking_exceptions:
|
|
481
|
+
# We emitted traced message but we didn't break on non_breaking_exception. We still need to raise an exception so that the
|
|
482
|
+
# call of `create_and_get_completed_partitions` knows that there was an issue with some partitions and the sync is incomplete.
|
|
483
|
+
raise AirbyteTracedException(
|
|
484
|
+
message="",
|
|
485
|
+
internal_message="\n".join(
|
|
486
|
+
[
|
|
487
|
+
filter_secrets(exception.__repr__())
|
|
488
|
+
for exception in self._non_breaking_exceptions
|
|
489
|
+
]
|
|
490
|
+
),
|
|
491
|
+
failure_type=FailureType.config_error,
|
|
492
|
+
)
|
|
493
|
+
|
|
494
|
+
def _handle_non_breaking_error(self, exception: Exception) -> None:
|
|
495
|
+
LOGGER.error(f"Failed to start the Job: {exception}, traceback: {traceback.format_exc()}")
|
|
496
|
+
self._non_breaking_exceptions.append(exception)
|
|
497
|
+
|
|
498
|
+
def _abort_all_running_jobs(self) -> None:
|
|
499
|
+
for partition in self._running_partitions:
|
|
500
|
+
for job in partition.jobs:
|
|
501
|
+
if job.status() in self._RUNNING_ON_API_SIDE_STATUS:
|
|
502
|
+
self._abort_job(job, free_job_allocation=True)
|
|
503
|
+
self._job_tracker.remove_job(job.api_job_id())
|
|
504
|
+
|
|
505
|
+
self._running_partitions = []
|
|
506
|
+
|
|
507
|
+
def _is_breaking_exception(self, exception: Exception) -> bool:
|
|
508
|
+
return isinstance(exception, self._exceptions_to_break_on) or (
|
|
509
|
+
isinstance(exception, AirbyteTracedException)
|
|
510
|
+
and exception.failure_type == FailureType.config_error
|
|
511
|
+
)
|
|
512
|
+
|
|
513
|
+
def fetch_records(self, async_jobs: Iterable[AsyncJob]) -> Iterable[Mapping[str, Any]]:
|
|
514
|
+
"""
|
|
515
|
+
Fetches records from the given jobs.
|
|
516
|
+
|
|
517
|
+
Args:
|
|
518
|
+
async_jobs Iterable[AsyncJob]: The list of AsyncJobs.
|
|
519
|
+
|
|
520
|
+
Yields:
|
|
521
|
+
Iterable[Mapping[str, Any]]: The fetched records from the jobs.
|
|
522
|
+
"""
|
|
523
|
+
for job in async_jobs:
|
|
524
|
+
yield from self._job_repository.fetch_records(job)
|
|
525
|
+
self._job_repository.delete(job)
|
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
|
|
2
|
+
|
|
3
|
+
import logging
|
|
4
|
+
import threading
|
|
5
|
+
import uuid
|
|
6
|
+
from typing import Set
|
|
7
|
+
|
|
8
|
+
from airbyte_cdk.logger import lazy_log
|
|
9
|
+
|
|
10
|
+
LOGGER = logging.getLogger("airbyte")
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class ConcurrentJobLimitReached(Exception):
|
|
14
|
+
pass
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class JobTracker:
|
|
18
|
+
def __init__(self, limit: int):
|
|
19
|
+
self._jobs: Set[str] = set()
|
|
20
|
+
if limit < 1:
|
|
21
|
+
LOGGER.warning(
|
|
22
|
+
f"The `max_concurrent_async_job_count` property is less than 1: {limit}. Setting to 1. Please update the source manifest to set a valid value."
|
|
23
|
+
)
|
|
24
|
+
self._limit = 1 if limit < 1 else limit
|
|
25
|
+
self._lock = threading.Lock()
|
|
26
|
+
|
|
27
|
+
def try_to_get_intent(self) -> str:
|
|
28
|
+
lazy_log(
|
|
29
|
+
LOGGER,
|
|
30
|
+
logging.DEBUG,
|
|
31
|
+
lambda: f"JobTracker - Trying to acquire lock by thread {threading.get_native_id()}...",
|
|
32
|
+
)
|
|
33
|
+
with self._lock:
|
|
34
|
+
if self._has_reached_limit():
|
|
35
|
+
raise ConcurrentJobLimitReached(
|
|
36
|
+
"Can't allocate more jobs right now: limit already reached"
|
|
37
|
+
)
|
|
38
|
+
intent = f"intent_{str(uuid.uuid4())}"
|
|
39
|
+
lazy_log(
|
|
40
|
+
LOGGER,
|
|
41
|
+
logging.DEBUG,
|
|
42
|
+
lambda: f"JobTracker - Thread {threading.get_native_id()} has acquired {intent}!",
|
|
43
|
+
)
|
|
44
|
+
self._jobs.add(intent)
|
|
45
|
+
return intent
|
|
46
|
+
|
|
47
|
+
def add_job(self, intent_or_job_id: str, job_id: str) -> None:
|
|
48
|
+
if intent_or_job_id not in self._jobs:
|
|
49
|
+
raise ValueError(
|
|
50
|
+
f"Can't add job: Unknown intent or job id, known values are {self._jobs}"
|
|
51
|
+
)
|
|
52
|
+
|
|
53
|
+
if intent_or_job_id == job_id:
|
|
54
|
+
# Nothing to do here as the ID to replace is the same
|
|
55
|
+
return
|
|
56
|
+
|
|
57
|
+
lazy_log(
|
|
58
|
+
LOGGER,
|
|
59
|
+
logging.DEBUG,
|
|
60
|
+
lambda: f"JobTracker - Thread {threading.get_native_id()} replacing job {intent_or_job_id} by {job_id}!",
|
|
61
|
+
)
|
|
62
|
+
with self._lock:
|
|
63
|
+
self._jobs.add(job_id)
|
|
64
|
+
self._jobs.remove(intent_or_job_id)
|
|
65
|
+
|
|
66
|
+
def remove_job(self, job_id: str) -> None:
|
|
67
|
+
"""
|
|
68
|
+
If the job is not allocated as a running job, this method does nothing and it won't raise.
|
|
69
|
+
"""
|
|
70
|
+
lazy_log(
|
|
71
|
+
LOGGER,
|
|
72
|
+
logging.DEBUG,
|
|
73
|
+
lambda: f"JobTracker - Thread {threading.get_native_id()} removing job {job_id}",
|
|
74
|
+
)
|
|
75
|
+
with self._lock:
|
|
76
|
+
self._jobs.discard(job_id)
|
|
77
|
+
|
|
78
|
+
def _has_reached_limit(self) -> bool:
|
|
79
|
+
return len(self._jobs) >= self._limit
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
|
|
2
|
+
|
|
3
|
+
from abc import abstractmethod
|
|
4
|
+
from typing import Any, Iterable, Mapping, Set
|
|
5
|
+
|
|
6
|
+
from airbyte_cdk.sources.declarative.async_job.job import AsyncJob
|
|
7
|
+
from airbyte_cdk.sources.types import StreamSlice
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class AsyncJobRepository:
|
|
11
|
+
@abstractmethod
|
|
12
|
+
def start(self, stream_slice: StreamSlice) -> AsyncJob:
|
|
13
|
+
pass
|
|
14
|
+
|
|
15
|
+
@abstractmethod
|
|
16
|
+
def update_jobs_status(self, jobs: Set[AsyncJob]) -> None:
|
|
17
|
+
pass
|
|
18
|
+
|
|
19
|
+
@abstractmethod
|
|
20
|
+
def fetch_records(self, job: AsyncJob) -> Iterable[Mapping[str, Any]]:
|
|
21
|
+
pass
|
|
22
|
+
|
|
23
|
+
@abstractmethod
|
|
24
|
+
def abort(self, job: AsyncJob) -> None:
|
|
25
|
+
"""
|
|
26
|
+
Called when we need to stop on the API side. This method can raise NotImplementedError as not all the APIs will support aborting
|
|
27
|
+
jobs.
|
|
28
|
+
"""
|
|
29
|
+
raise NotImplementedError(
|
|
30
|
+
"Either the API or the AsyncJobRepository implementation do not support aborting jobs"
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
@abstractmethod
|
|
34
|
+
def delete(self, job: AsyncJob) -> None:
|
|
35
|
+
pass
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
from enum import Enum
|
|
5
|
+
|
|
6
|
+
_TERMINAL = True
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class AsyncJobStatus(Enum):
|
|
10
|
+
RUNNING = ("RUNNING", not _TERMINAL)
|
|
11
|
+
COMPLETED = ("COMPLETED", _TERMINAL)
|
|
12
|
+
FAILED = ("FAILED", _TERMINAL)
|
|
13
|
+
TIMED_OUT = ("TIMED_OUT", _TERMINAL)
|
|
14
|
+
|
|
15
|
+
def __init__(self, value: str, is_terminal: bool) -> None:
|
|
16
|
+
self._value = value
|
|
17
|
+
self._is_terminal = is_terminal
|
|
18
|
+
|
|
19
|
+
def is_terminal(self) -> bool:
|
|
20
|
+
"""
|
|
21
|
+
A status is terminal when a job status can't be updated anymore. For example if a job is completed, it will stay completed but a
|
|
22
|
+
running job might because completed, failed or timed out.
|
|
23
|
+
"""
|
|
24
|
+
return self._is_terminal
|