airbyte-cdk 0.72.1__py3-none-any.whl → 6.13.1.dev4106__py3-none-any.whl
Sign up to get free protection for your applications and to get access to all the features.
- airbyte_cdk/__init__.py +355 -6
- 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 +230 -0
- airbyte_cdk/cli/source_declarative_manifest/spec.json +17 -0
- airbyte_cdk/config_observation.py +29 -10
- airbyte_cdk/connector.py +24 -24
- airbyte_cdk/connector_builder/README.md +53 -0
- airbyte_cdk/connector_builder/connector_builder_handler.py +37 -11
- airbyte_cdk/connector_builder/main.py +45 -13
- airbyte_cdk/connector_builder/message_grouper.py +189 -50
- airbyte_cdk/connector_builder/models.py +3 -2
- airbyte_cdk/destinations/__init__.py +4 -3
- airbyte_cdk/destinations/destination.py +54 -20
- airbyte_cdk/destinations/vector_db_based/README.md +37 -0
- airbyte_cdk/destinations/vector_db_based/config.py +40 -17
- airbyte_cdk/destinations/vector_db_based/document_processor.py +56 -17
- airbyte_cdk/destinations/vector_db_based/embedder.py +57 -15
- airbyte_cdk/destinations/vector_db_based/test_utils.py +14 -4
- airbyte_cdk/destinations/vector_db_based/utils.py +8 -2
- airbyte_cdk/destinations/vector_db_based/writer.py +24 -5
- airbyte_cdk/entrypoint.py +153 -44
- airbyte_cdk/exception_handler.py +21 -3
- airbyte_cdk/logger.py +30 -44
- airbyte_cdk/models/__init__.py +13 -2
- airbyte_cdk/models/airbyte_protocol.py +86 -1
- airbyte_cdk/models/airbyte_protocol_serializers.py +44 -0
- airbyte_cdk/models/file_transfer_record_message.py +13 -0
- airbyte_cdk/models/well_known_types.py +1 -1
- airbyte_cdk/sources/__init__.py +5 -1
- airbyte_cdk/sources/abstract_source.py +125 -79
- airbyte_cdk/sources/concurrent_source/__init__.py +7 -2
- airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py +102 -36
- airbyte_cdk/sources/concurrent_source/concurrent_source.py +29 -36
- airbyte_cdk/sources/concurrent_source/concurrent_source_adapter.py +94 -10
- airbyte_cdk/sources/concurrent_source/stream_thread_exception.py +25 -0
- airbyte_cdk/sources/concurrent_source/thread_pool_manager.py +20 -14
- airbyte_cdk/sources/config.py +3 -2
- airbyte_cdk/sources/connector_state_manager.py +49 -83
- airbyte_cdk/sources/declarative/async_job/job.py +52 -0
- airbyte_cdk/sources/declarative/async_job/job_orchestrator.py +497 -0
- airbyte_cdk/sources/declarative/async_job/job_tracker.py +75 -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 +2 -3
- airbyte_cdk/sources/declarative/auth/declarative_authenticator.py +3 -1
- airbyte_cdk/sources/declarative/auth/jwt.py +191 -0
- airbyte_cdk/sources/declarative/auth/oauth.py +60 -20
- airbyte_cdk/sources/declarative/auth/selective_authenticator.py +10 -2
- airbyte_cdk/sources/declarative/auth/token.py +28 -10
- airbyte_cdk/sources/declarative/auth/token_provider.py +9 -8
- airbyte_cdk/sources/declarative/checks/check_stream.py +16 -8
- airbyte_cdk/sources/declarative/checks/connection_checker.py +4 -2
- 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 +421 -0
- airbyte_cdk/sources/declarative/datetime/datetime_parser.py +4 -0
- airbyte_cdk/sources/declarative/datetime/min_max_datetime.py +26 -6
- airbyte_cdk/sources/declarative/declarative_component_schema.yaml +1185 -85
- airbyte_cdk/sources/declarative/declarative_source.py +5 -2
- airbyte_cdk/sources/declarative/declarative_stream.py +95 -9
- airbyte_cdk/sources/declarative/decoders/__init__.py +23 -2
- airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py +97 -0
- airbyte_cdk/sources/declarative/decoders/decoder.py +11 -4
- airbyte_cdk/sources/declarative/decoders/json_decoder.py +92 -5
- 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/extractors/__init__.py +12 -1
- airbyte_cdk/sources/declarative/extractors/dpath_extractor.py +29 -24
- airbyte_cdk/sources/declarative/extractors/http_selector.py +4 -5
- airbyte_cdk/sources/declarative/extractors/record_extractor.py +2 -3
- airbyte_cdk/sources/declarative/extractors/record_filter.py +65 -8
- airbyte_cdk/sources/declarative/extractors/record_selector.py +85 -26
- airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py +177 -0
- airbyte_cdk/sources/declarative/extractors/type_transformer.py +55 -0
- airbyte_cdk/sources/declarative/incremental/__init__.py +25 -3
- airbyte_cdk/sources/declarative/incremental/datetime_based_cursor.py +156 -48
- airbyte_cdk/sources/declarative/incremental/declarative_cursor.py +13 -0
- airbyte_cdk/sources/declarative/incremental/global_substream_cursor.py +350 -0
- airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py +159 -74
- 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/filters.py +27 -1
- airbyte_cdk/sources/declarative/interpolation/interpolated_boolean.py +23 -5
- airbyte_cdk/sources/declarative/interpolation/interpolated_mapping.py +12 -8
- airbyte_cdk/sources/declarative/interpolation/interpolated_nested_mapping.py +13 -6
- airbyte_cdk/sources/declarative/interpolation/interpolated_string.py +21 -6
- airbyte_cdk/sources/declarative/interpolation/interpolation.py +9 -3
- airbyte_cdk/sources/declarative/interpolation/jinja.py +72 -37
- airbyte_cdk/sources/declarative/interpolation/macros.py +72 -17
- airbyte_cdk/sources/declarative/manifest_declarative_source.py +193 -52
- 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 +1 -1
- airbyte_cdk/sources/declarative/models/declarative_component_schema.py +1319 -603
- airbyte_cdk/sources/declarative/parsers/custom_exceptions.py +2 -2
- airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py +26 -4
- airbyte_cdk/sources/declarative/parsers/manifest_reference_resolver.py +26 -15
- airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +1695 -225
- airbyte_cdk/sources/declarative/partition_routers/__init__.py +24 -4
- 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 +39 -9
- airbyte_cdk/sources/declarative/partition_routers/partition_router.py +62 -0
- airbyte_cdk/sources/declarative/partition_routers/single_partition_router.py +15 -3
- airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +222 -39
- airbyte_cdk/sources/declarative/requesters/error_handlers/__init__.py +19 -5
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/__init__.py +3 -1
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/constant_backoff_strategy.py +19 -7
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/exponential_backoff_strategy.py +19 -7
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/header_helper.py +4 -2
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_time_from_header_backoff_strategy.py +41 -9
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_until_time_from_header_backoff_strategy.py +29 -14
- airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategy.py +5 -13
- airbyte_cdk/sources/declarative/requesters/error_handlers/composite_error_handler.py +32 -16
- airbyte_cdk/sources/declarative/requesters/error_handlers/default_error_handler.py +46 -56
- airbyte_cdk/sources/declarative/requesters/error_handlers/default_http_response_filter.py +40 -0
- airbyte_cdk/sources/declarative/requesters/error_handlers/error_handler.py +6 -32
- airbyte_cdk/sources/declarative/requesters/error_handlers/http_response_filter.py +119 -41
- airbyte_cdk/sources/declarative/requesters/http_job_repository.py +228 -0
- airbyte_cdk/sources/declarative/requesters/http_requester.py +98 -344
- airbyte_cdk/sources/declarative/requesters/paginators/__init__.py +14 -3
- airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py +105 -46
- airbyte_cdk/sources/declarative/requesters/paginators/no_pagination.py +14 -8
- airbyte_cdk/sources/declarative/requesters/paginators/paginator.py +19 -8
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/__init__.py +9 -3
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/cursor_pagination_strategy.py +53 -21
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py +42 -19
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/page_increment.py +25 -12
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py +13 -10
- airbyte_cdk/sources/declarative/requesters/paginators/strategies/stop_condition.py +26 -13
- airbyte_cdk/sources/declarative/requesters/request_options/__init__.py +15 -2
- airbyte_cdk/sources/declarative/requesters/request_options/datetime_based_request_options_provider.py +91 -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 +31 -14
- airbyte_cdk/sources/declarative/requesters/request_options/interpolated_request_input_provider.py +27 -15
- airbyte_cdk/sources/declarative/requesters/request_options/interpolated_request_options_provider.py +63 -10
- airbyte_cdk/sources/declarative/requesters/request_options/request_options_provider.py +1 -1
- airbyte_cdk/sources/declarative/requesters/requester.py +9 -17
- 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 +6 -2
- airbyte_cdk/sources/declarative/retrievers/async_retriever.py +100 -0
- airbyte_cdk/sources/declarative/retrievers/retriever.py +1 -3
- airbyte_cdk/sources/declarative/retrievers/simple_retriever.py +228 -72
- airbyte_cdk/sources/declarative/schema/__init__.py +14 -1
- airbyte_cdk/sources/declarative/schema/default_schema_loader.py +5 -3
- airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py +236 -0
- airbyte_cdk/sources/declarative/schema/json_file_schema_loader.py +8 -8
- airbyte_cdk/sources/declarative/spec/spec.py +12 -5
- airbyte_cdk/sources/declarative/stream_slicers/__init__.py +1 -2
- airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py +88 -0
- airbyte_cdk/sources/declarative/stream_slicers/stream_slicer.py +9 -14
- airbyte_cdk/sources/declarative/transformations/add_fields.py +19 -11
- 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 +13 -10
- airbyte_cdk/sources/declarative/transformations/transformation.py +5 -5
- airbyte_cdk/sources/declarative/types.py +19 -110
- airbyte_cdk/sources/declarative/yaml_declarative_source.py +31 -10
- airbyte_cdk/sources/embedded/base_integration.py +16 -5
- airbyte_cdk/sources/embedded/catalog.py +16 -4
- airbyte_cdk/sources/embedded/runner.py +19 -3
- airbyte_cdk/sources/embedded/tools.py +5 -2
- 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 +9 -2
- airbyte_cdk/sources/file_based/availability_strategy/abstract_file_based_availability_strategy.py +22 -6
- airbyte_cdk/sources/file_based/availability_strategy/default_file_based_availability_strategy.py +46 -10
- airbyte_cdk/sources/file_based/config/abstract_file_based_spec.py +58 -10
- airbyte_cdk/sources/file_based/config/avro_format.py +2 -1
- airbyte_cdk/sources/file_based/config/csv_format.py +29 -10
- airbyte_cdk/sources/file_based/config/excel_format.py +18 -0
- airbyte_cdk/sources/file_based/config/file_based_stream_config.py +16 -4
- airbyte_cdk/sources/file_based/config/jsonl_format.py +2 -1
- airbyte_cdk/sources/file_based/config/parquet_format.py +2 -1
- airbyte_cdk/sources/file_based/config/unstructured_format.py +13 -5
- airbyte_cdk/sources/file_based/discovery_policy/__init__.py +6 -2
- airbyte_cdk/sources/file_based/discovery_policy/abstract_discovery_policy.py +2 -4
- airbyte_cdk/sources/file_based/discovery_policy/default_discovery_policy.py +7 -2
- airbyte_cdk/sources/file_based/exceptions.py +52 -15
- airbyte_cdk/sources/file_based/file_based_source.py +163 -33
- airbyte_cdk/sources/file_based/file_based_stream_reader.py +83 -5
- airbyte_cdk/sources/file_based/file_types/__init__.py +14 -1
- airbyte_cdk/sources/file_based/file_types/avro_parser.py +75 -24
- airbyte_cdk/sources/file_based/file_types/csv_parser.py +116 -34
- airbyte_cdk/sources/file_based/file_types/excel_parser.py +196 -0
- airbyte_cdk/sources/file_based/file_types/file_transfer.py +37 -0
- airbyte_cdk/sources/file_based/file_types/file_type_parser.py +4 -1
- airbyte_cdk/sources/file_based/file_types/jsonl_parser.py +24 -8
- airbyte_cdk/sources/file_based/file_types/parquet_parser.py +60 -18
- airbyte_cdk/sources/file_based/file_types/unstructured_parser.py +145 -41
- airbyte_cdk/sources/file_based/remote_file.py +1 -1
- airbyte_cdk/sources/file_based/schema_helpers.py +38 -10
- airbyte_cdk/sources/file_based/schema_validation_policies/__init__.py +3 -1
- airbyte_cdk/sources/file_based/schema_validation_policies/abstract_schema_validation_policy.py +3 -1
- airbyte_cdk/sources/file_based/schema_validation_policies/default_schema_validation_policies.py +16 -5
- airbyte_cdk/sources/file_based/stream/abstract_file_based_stream.py +50 -13
- airbyte_cdk/sources/file_based/stream/concurrent/adapters.py +67 -27
- airbyte_cdk/sources/file_based/stream/concurrent/cursor/__init__.py +5 -1
- airbyte_cdk/sources/file_based/stream/concurrent/cursor/abstract_concurrent_file_based_cursor.py +14 -23
- airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_concurrent_cursor.py +54 -18
- airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_final_state_cursor.py +21 -9
- airbyte_cdk/sources/file_based/stream/cursor/abstract_file_based_cursor.py +3 -1
- airbyte_cdk/sources/file_based/stream/cursor/default_file_based_cursor.py +27 -10
- airbyte_cdk/sources/file_based/stream/default_file_based_stream.py +175 -45
- airbyte_cdk/sources/http_logger.py +8 -3
- airbyte_cdk/sources/message/__init__.py +7 -1
- airbyte_cdk/sources/message/repository.py +18 -4
- airbyte_cdk/sources/source.py +42 -38
- airbyte_cdk/sources/streams/__init__.py +2 -2
- airbyte_cdk/sources/streams/availability_strategy.py +54 -3
- airbyte_cdk/sources/streams/call_rate.py +64 -21
- airbyte_cdk/sources/streams/checkpoint/__init__.py +26 -0
- airbyte_cdk/sources/streams/checkpoint/checkpoint_reader.py +335 -0
- airbyte_cdk/sources/{declarative/incremental → streams/checkpoint}/cursor.py +17 -14
- 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/abstract_stream.py +7 -2
- airbyte_cdk/sources/streams/concurrent/adapters.py +84 -75
- airbyte_cdk/sources/streams/concurrent/availability_strategy.py +30 -2
- airbyte_cdk/sources/streams/concurrent/cursor.py +298 -42
- airbyte_cdk/sources/streams/concurrent/default_stream.py +12 -3
- airbyte_cdk/sources/streams/concurrent/exceptions.py +3 -0
- airbyte_cdk/sources/streams/concurrent/helpers.py +14 -3
- airbyte_cdk/sources/streams/concurrent/partition_enqueuer.py +12 -3
- airbyte_cdk/sources/streams/concurrent/partition_reader.py +10 -3
- airbyte_cdk/sources/streams/concurrent/partitions/partition.py +1 -16
- airbyte_cdk/sources/streams/concurrent/partitions/stream_slicer.py +21 -0
- airbyte_cdk/sources/streams/concurrent/partitions/types.py +15 -5
- airbyte_cdk/sources/streams/concurrent/state_converters/abstract_stream_state_converter.py +109 -17
- airbyte_cdk/sources/streams/concurrent/state_converters/datetime_stream_state_converter.py +90 -72
- airbyte_cdk/sources/streams/core.py +412 -87
- airbyte_cdk/sources/streams/http/__init__.py +2 -1
- airbyte_cdk/sources/streams/http/availability_strategy.py +12 -101
- 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 +27 -7
- airbyte_cdk/sources/streams/http/http.py +369 -246
- airbyte_cdk/sources/streams/http/http_client.py +531 -0
- airbyte_cdk/sources/streams/http/rate_limiting.py +76 -12
- airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py +28 -9
- airbyte_cdk/sources/streams/http/requests_native_auth/abstract_token.py +2 -1
- airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py +90 -35
- airbyte_cdk/sources/streams/http/requests_native_auth/token.py +13 -3
- airbyte_cdk/sources/types.py +154 -0
- airbyte_cdk/sources/utils/record_helper.py +36 -21
- airbyte_cdk/sources/utils/schema_helpers.py +13 -6
- airbyte_cdk/sources/utils/slice_logger.py +4 -1
- airbyte_cdk/sources/utils/transform.py +54 -20
- 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/catalog_builder.py +70 -18
- airbyte_cdk/test/entrypoint_wrapper.py +117 -42
- airbyte_cdk/test/mock_http/__init__.py +1 -1
- airbyte_cdk/test/mock_http/matcher.py +6 -0
- airbyte_cdk/test/mock_http/mocker.py +57 -10
- airbyte_cdk/test/mock_http/request.py +19 -3
- airbyte_cdk/test/mock_http/response.py +3 -1
- airbyte_cdk/test/mock_http/response_builder.py +32 -16
- airbyte_cdk/test/state_builder.py +18 -10
- 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 +60 -0
- airbyte_cdk/test/utils/reading.py +26 -0
- airbyte_cdk/utils/__init__.py +2 -1
- airbyte_cdk/utils/airbyte_secrets_utils.py +5 -3
- airbyte_cdk/utils/analytics_message.py +10 -2
- airbyte_cdk/utils/datetime_format_inferrer.py +4 -1
- airbyte_cdk/utils/event_timing.py +10 -10
- airbyte_cdk/utils/mapping_helpers.py +3 -1
- airbyte_cdk/utils/message_utils.py +20 -11
- airbyte_cdk/utils/print_buffer.py +75 -0
- airbyte_cdk/utils/schema_inferrer.py +198 -28
- airbyte_cdk/utils/slice_hasher.py +30 -0
- airbyte_cdk/utils/spec_schema_transformations.py +6 -3
- airbyte_cdk/utils/stream_status_utils.py +8 -1
- airbyte_cdk/utils/traced_exception.py +61 -21
- airbyte_cdk-6.13.1.dev4106.dist-info/METADATA +109 -0
- airbyte_cdk-6.13.1.dev4106.dist-info/RECORD +349 -0
- {airbyte_cdk-0.72.1.dist-info → airbyte_cdk-6.13.1.dev4106.dist-info}/WHEEL +1 -2
- airbyte_cdk-6.13.1.dev4106.dist-info/entry_points.txt +3 -0
- airbyte_cdk/sources/declarative/create_partial.py +0 -92
- airbyte_cdk/sources/declarative/parsers/class_types_registry.py +0 -102
- airbyte_cdk/sources/declarative/parsers/default_implementation_registry.py +0 -64
- airbyte_cdk/sources/declarative/requesters/error_handlers/response_action.py +0 -16
- airbyte_cdk/sources/declarative/requesters/error_handlers/response_status.py +0 -68
- airbyte_cdk/sources/declarative/stream_slicers/cartesian_product_stream_slicer.py +0 -114
- airbyte_cdk/sources/deprecated/base_source.py +0 -94
- airbyte_cdk/sources/deprecated/client.py +0 -99
- airbyte_cdk/sources/singer/__init__.py +0 -8
- airbyte_cdk/sources/singer/singer_helpers.py +0 -304
- airbyte_cdk/sources/singer/source.py +0 -186
- airbyte_cdk/sources/streams/concurrent/partitions/record.py +0 -23
- airbyte_cdk/sources/streams/http/auth/__init__.py +0 -17
- airbyte_cdk/sources/streams/http/auth/core.py +0 -29
- airbyte_cdk/sources/streams/http/auth/oauth.py +0 -113
- airbyte_cdk/sources/streams/http/auth/token.py +0 -47
- airbyte_cdk/sources/streams/utils/stream_helper.py +0 -40
- airbyte_cdk/sources/utils/catalog_helpers.py +0 -22
- airbyte_cdk/sources/utils/schema_models.py +0 -84
- airbyte_cdk-0.72.1.dist-info/METADATA +0 -243
- airbyte_cdk-0.72.1.dist-info/RECORD +0 -466
- airbyte_cdk-0.72.1.dist-info/top_level.txt +0 -3
- source_declarative_manifest/main.py +0 -29
- unit_tests/connector_builder/__init__.py +0 -3
- unit_tests/connector_builder/test_connector_builder_handler.py +0 -871
- unit_tests/connector_builder/test_message_grouper.py +0 -713
- unit_tests/connector_builder/utils.py +0 -27
- unit_tests/destinations/test_destination.py +0 -243
- unit_tests/singer/test_singer_helpers.py +0 -56
- unit_tests/singer/test_singer_source.py +0 -112
- unit_tests/sources/__init__.py +0 -0
- unit_tests/sources/concurrent_source/__init__.py +0 -3
- unit_tests/sources/concurrent_source/test_concurrent_source_adapter.py +0 -106
- unit_tests/sources/declarative/__init__.py +0 -3
- unit_tests/sources/declarative/auth/__init__.py +0 -3
- unit_tests/sources/declarative/auth/test_oauth.py +0 -331
- unit_tests/sources/declarative/auth/test_selective_authenticator.py +0 -39
- unit_tests/sources/declarative/auth/test_session_token_auth.py +0 -182
- unit_tests/sources/declarative/auth/test_token_auth.py +0 -200
- unit_tests/sources/declarative/auth/test_token_provider.py +0 -73
- unit_tests/sources/declarative/checks/__init__.py +0 -3
- unit_tests/sources/declarative/checks/test_check_stream.py +0 -146
- unit_tests/sources/declarative/decoders/__init__.py +0 -0
- unit_tests/sources/declarative/decoders/test_json_decoder.py +0 -16
- unit_tests/sources/declarative/external_component.py +0 -13
- unit_tests/sources/declarative/extractors/__init__.py +0 -3
- unit_tests/sources/declarative/extractors/test_dpath_extractor.py +0 -55
- unit_tests/sources/declarative/extractors/test_record_filter.py +0 -55
- unit_tests/sources/declarative/extractors/test_record_selector.py +0 -179
- unit_tests/sources/declarative/incremental/__init__.py +0 -0
- unit_tests/sources/declarative/incremental/test_datetime_based_cursor.py +0 -860
- unit_tests/sources/declarative/incremental/test_per_partition_cursor.py +0 -406
- unit_tests/sources/declarative/incremental/test_per_partition_cursor_integration.py +0 -332
- unit_tests/sources/declarative/interpolation/__init__.py +0 -3
- unit_tests/sources/declarative/interpolation/test_filters.py +0 -80
- unit_tests/sources/declarative/interpolation/test_interpolated_boolean.py +0 -40
- unit_tests/sources/declarative/interpolation/test_interpolated_mapping.py +0 -35
- unit_tests/sources/declarative/interpolation/test_interpolated_nested_mapping.py +0 -45
- unit_tests/sources/declarative/interpolation/test_interpolated_string.py +0 -25
- unit_tests/sources/declarative/interpolation/test_jinja.py +0 -240
- unit_tests/sources/declarative/interpolation/test_macros.py +0 -73
- unit_tests/sources/declarative/parsers/__init__.py +0 -3
- unit_tests/sources/declarative/parsers/test_manifest_component_transformer.py +0 -406
- unit_tests/sources/declarative/parsers/test_manifest_reference_resolver.py +0 -139
- unit_tests/sources/declarative/parsers/test_model_to_component_factory.py +0 -1847
- unit_tests/sources/declarative/parsers/testing_components.py +0 -36
- unit_tests/sources/declarative/partition_routers/__init__.py +0 -3
- unit_tests/sources/declarative/partition_routers/test_list_partition_router.py +0 -155
- unit_tests/sources/declarative/partition_routers/test_single_partition_router.py +0 -14
- unit_tests/sources/declarative/partition_routers/test_substream_partition_router.py +0 -404
- unit_tests/sources/declarative/requesters/__init__.py +0 -3
- unit_tests/sources/declarative/requesters/error_handlers/__init__.py +0 -3
- unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/__init__.py +0 -3
- unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_constant_backoff.py +0 -34
- unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_exponential_backoff.py +0 -36
- unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_header_helper.py +0 -38
- unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_wait_time_from_header.py +0 -35
- unit_tests/sources/declarative/requesters/error_handlers/backoff_strategies/test_wait_until_time_from_header.py +0 -64
- unit_tests/sources/declarative/requesters/error_handlers/test_composite_error_handler.py +0 -213
- unit_tests/sources/declarative/requesters/error_handlers/test_default_error_handler.py +0 -178
- unit_tests/sources/declarative/requesters/error_handlers/test_http_response_filter.py +0 -121
- unit_tests/sources/declarative/requesters/error_handlers/test_response_status.py +0 -44
- unit_tests/sources/declarative/requesters/paginators/__init__.py +0 -3
- unit_tests/sources/declarative/requesters/paginators/test_cursor_pagination_strategy.py +0 -64
- unit_tests/sources/declarative/requesters/paginators/test_default_paginator.py +0 -313
- unit_tests/sources/declarative/requesters/paginators/test_no_paginator.py +0 -12
- unit_tests/sources/declarative/requesters/paginators/test_offset_increment.py +0 -58
- unit_tests/sources/declarative/requesters/paginators/test_page_increment.py +0 -70
- unit_tests/sources/declarative/requesters/paginators/test_request_option.py +0 -43
- unit_tests/sources/declarative/requesters/paginators/test_stop_condition.py +0 -105
- unit_tests/sources/declarative/requesters/request_options/__init__.py +0 -3
- unit_tests/sources/declarative/requesters/request_options/test_interpolated_request_options_provider.py +0 -101
- unit_tests/sources/declarative/requesters/test_http_requester.py +0 -974
- unit_tests/sources/declarative/requesters/test_interpolated_request_input_provider.py +0 -32
- unit_tests/sources/declarative/retrievers/__init__.py +0 -3
- unit_tests/sources/declarative/retrievers/test_simple_retriever.py +0 -542
- unit_tests/sources/declarative/schema/__init__.py +0 -6
- unit_tests/sources/declarative/schema/source_test/SourceTest.py +0 -8
- unit_tests/sources/declarative/schema/source_test/__init__.py +0 -3
- unit_tests/sources/declarative/schema/test_default_schema_loader.py +0 -32
- unit_tests/sources/declarative/schema/test_inline_schema_loader.py +0 -19
- unit_tests/sources/declarative/schema/test_json_file_schema_loader.py +0 -26
- unit_tests/sources/declarative/states/__init__.py +0 -3
- unit_tests/sources/declarative/stream_slicers/__init__.py +0 -3
- unit_tests/sources/declarative/stream_slicers/test_cartesian_product_stream_slicer.py +0 -225
- unit_tests/sources/declarative/test_create_partial.py +0 -83
- unit_tests/sources/declarative/test_declarative_stream.py +0 -103
- unit_tests/sources/declarative/test_manifest_declarative_source.py +0 -1260
- unit_tests/sources/declarative/test_types.py +0 -39
- unit_tests/sources/declarative/test_yaml_declarative_source.py +0 -148
- unit_tests/sources/file_based/__init__.py +0 -0
- unit_tests/sources/file_based/availability_strategy/__init__.py +0 -0
- unit_tests/sources/file_based/availability_strategy/test_default_file_based_availability_strategy.py +0 -100
- unit_tests/sources/file_based/config/__init__.py +0 -0
- unit_tests/sources/file_based/config/test_abstract_file_based_spec.py +0 -28
- unit_tests/sources/file_based/config/test_csv_format.py +0 -34
- unit_tests/sources/file_based/config/test_file_based_stream_config.py +0 -84
- unit_tests/sources/file_based/discovery_policy/__init__.py +0 -0
- unit_tests/sources/file_based/discovery_policy/test_default_discovery_policy.py +0 -31
- unit_tests/sources/file_based/file_types/__init__.py +0 -0
- unit_tests/sources/file_based/file_types/test_avro_parser.py +0 -243
- unit_tests/sources/file_based/file_types/test_csv_parser.py +0 -546
- unit_tests/sources/file_based/file_types/test_jsonl_parser.py +0 -158
- unit_tests/sources/file_based/file_types/test_parquet_parser.py +0 -274
- unit_tests/sources/file_based/file_types/test_unstructured_parser.py +0 -593
- unit_tests/sources/file_based/helpers.py +0 -70
- unit_tests/sources/file_based/in_memory_files_source.py +0 -211
- unit_tests/sources/file_based/scenarios/__init__.py +0 -0
- unit_tests/sources/file_based/scenarios/avro_scenarios.py +0 -744
- unit_tests/sources/file_based/scenarios/check_scenarios.py +0 -220
- unit_tests/sources/file_based/scenarios/concurrent_incremental_scenarios.py +0 -2844
- unit_tests/sources/file_based/scenarios/csv_scenarios.py +0 -3105
- unit_tests/sources/file_based/scenarios/file_based_source_builder.py +0 -91
- unit_tests/sources/file_based/scenarios/incremental_scenarios.py +0 -1926
- unit_tests/sources/file_based/scenarios/jsonl_scenarios.py +0 -930
- unit_tests/sources/file_based/scenarios/parquet_scenarios.py +0 -754
- unit_tests/sources/file_based/scenarios/scenario_builder.py +0 -234
- unit_tests/sources/file_based/scenarios/unstructured_scenarios.py +0 -608
- unit_tests/sources/file_based/scenarios/user_input_schema_scenarios.py +0 -746
- unit_tests/sources/file_based/scenarios/validation_policy_scenarios.py +0 -726
- unit_tests/sources/file_based/stream/__init__.py +0 -0
- unit_tests/sources/file_based/stream/concurrent/__init__.py +0 -0
- unit_tests/sources/file_based/stream/concurrent/test_adapters.py +0 -362
- unit_tests/sources/file_based/stream/concurrent/test_file_based_concurrent_cursor.py +0 -458
- unit_tests/sources/file_based/stream/test_default_file_based_cursor.py +0 -310
- unit_tests/sources/file_based/stream/test_default_file_based_stream.py +0 -244
- unit_tests/sources/file_based/test_file_based_scenarios.py +0 -320
- unit_tests/sources/file_based/test_file_based_stream_reader.py +0 -272
- unit_tests/sources/file_based/test_scenarios.py +0 -253
- unit_tests/sources/file_based/test_schema_helpers.py +0 -346
- unit_tests/sources/fixtures/__init__.py +0 -3
- unit_tests/sources/fixtures/source_test_fixture.py +0 -153
- unit_tests/sources/message/__init__.py +0 -0
- unit_tests/sources/message/test_repository.py +0 -153
- unit_tests/sources/streams/__init__.py +0 -0
- unit_tests/sources/streams/concurrent/__init__.py +0 -3
- unit_tests/sources/streams/concurrent/scenarios/__init__.py +0 -3
- unit_tests/sources/streams/concurrent/scenarios/incremental_scenarios.py +0 -250
- unit_tests/sources/streams/concurrent/scenarios/stream_facade_builder.py +0 -140
- unit_tests/sources/streams/concurrent/scenarios/stream_facade_scenarios.py +0 -452
- unit_tests/sources/streams/concurrent/scenarios/test_concurrent_scenarios.py +0 -76
- unit_tests/sources/streams/concurrent/scenarios/thread_based_concurrent_stream_scenarios.py +0 -418
- unit_tests/sources/streams/concurrent/scenarios/thread_based_concurrent_stream_source_builder.py +0 -142
- unit_tests/sources/streams/concurrent/scenarios/utils.py +0 -55
- unit_tests/sources/streams/concurrent/test_adapters.py +0 -380
- unit_tests/sources/streams/concurrent/test_concurrent_read_processor.py +0 -684
- unit_tests/sources/streams/concurrent/test_cursor.py +0 -139
- unit_tests/sources/streams/concurrent/test_datetime_state_converter.py +0 -369
- unit_tests/sources/streams/concurrent/test_default_stream.py +0 -197
- unit_tests/sources/streams/concurrent/test_partition_enqueuer.py +0 -90
- unit_tests/sources/streams/concurrent/test_partition_reader.py +0 -67
- unit_tests/sources/streams/concurrent/test_thread_pool_manager.py +0 -106
- unit_tests/sources/streams/http/__init__.py +0 -0
- unit_tests/sources/streams/http/auth/__init__.py +0 -0
- unit_tests/sources/streams/http/auth/test_auth.py +0 -173
- unit_tests/sources/streams/http/requests_native_auth/__init__.py +0 -0
- unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py +0 -423
- unit_tests/sources/streams/http/test_availability_strategy.py +0 -180
- unit_tests/sources/streams/http/test_http.py +0 -635
- unit_tests/sources/streams/test_availability_strategy.py +0 -70
- unit_tests/sources/streams/test_call_rate.py +0 -300
- unit_tests/sources/streams/test_stream_read.py +0 -405
- unit_tests/sources/streams/test_streams_core.py +0 -184
- unit_tests/sources/test_abstract_source.py +0 -1442
- unit_tests/sources/test_concurrent_source.py +0 -112
- unit_tests/sources/test_config.py +0 -92
- unit_tests/sources/test_connector_state_manager.py +0 -482
- unit_tests/sources/test_http_logger.py +0 -252
- unit_tests/sources/test_integration_source.py +0 -86
- unit_tests/sources/test_source.py +0 -684
- unit_tests/sources/test_source_read.py +0 -460
- unit_tests/test/__init__.py +0 -0
- unit_tests/test/mock_http/__init__.py +0 -0
- unit_tests/test/mock_http/test_matcher.py +0 -53
- unit_tests/test/mock_http/test_mocker.py +0 -214
- unit_tests/test/mock_http/test_request.py +0 -117
- unit_tests/test/mock_http/test_response_builder.py +0 -177
- unit_tests/test/test_entrypoint_wrapper.py +0 -240
- unit_tests/utils/__init__.py +0 -0
- unit_tests/utils/test_datetime_format_inferrer.py +0 -60
- unit_tests/utils/test_mapping_helpers.py +0 -54
- unit_tests/utils/test_message_utils.py +0 -91
- unit_tests/utils/test_rate_limiting.py +0 -26
- unit_tests/utils/test_schema_inferrer.py +0 -202
- unit_tests/utils/test_secret_utils.py +0 -135
- unit_tests/utils/test_stream_status_utils.py +0 -61
- unit_tests/utils/test_traced_exception.py +0 -107
- /airbyte_cdk/sources/{deprecated → declarative/async_job}/__init__.py +0 -0
- {source_declarative_manifest → airbyte_cdk/sources/declarative/migrations}/__init__.py +0 -0
- {unit_tests/destinations → airbyte_cdk/sql}/__init__.py +0 -0
- {unit_tests/singer → airbyte_cdk/sql/_util}/__init__.py +0 -0
- {airbyte_cdk-0.72.1.dist-info → airbyte_cdk-6.13.1.dev4106.dist-info}/LICENSE.txt +0 -0
@@ -0,0 +1,786 @@
|
|
1
|
+
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
|
2
|
+
"""The base SQL Cache implementation."""
|
3
|
+
|
4
|
+
from __future__ import annotations
|
5
|
+
|
6
|
+
import abc
|
7
|
+
from collections import defaultdict
|
8
|
+
from contextlib import contextmanager
|
9
|
+
from functools import cached_property
|
10
|
+
from pathlib import Path
|
11
|
+
from typing import TYPE_CHECKING, Any, final
|
12
|
+
|
13
|
+
import pandas as pd
|
14
|
+
import sqlalchemy
|
15
|
+
import ulid
|
16
|
+
from airbyte_protocol_dataclasses.models import AirbyteStateMessage
|
17
|
+
from pandas import Index
|
18
|
+
from pydantic import BaseModel, Field
|
19
|
+
from sqlalchemy import Column, Table, and_, create_engine, insert, null, select, text, update
|
20
|
+
from sqlalchemy.exc import ProgrammingError, SQLAlchemyError
|
21
|
+
|
22
|
+
from airbyte_cdk.sql import exceptions as exc
|
23
|
+
from airbyte_cdk.sql._util.hashing import one_way_hash
|
24
|
+
from airbyte_cdk.sql._util.name_normalizers import LowerCaseNormalizer
|
25
|
+
from airbyte_cdk.sql.constants import (
|
26
|
+
AB_EXTRACTED_AT_COLUMN,
|
27
|
+
AB_META_COLUMN,
|
28
|
+
AB_RAW_ID_COLUMN,
|
29
|
+
DEBUG_MODE,
|
30
|
+
)
|
31
|
+
from airbyte_cdk.sql.secrets import SecretString
|
32
|
+
from airbyte_cdk.sql.types import SQLTypeConverter
|
33
|
+
|
34
|
+
if TYPE_CHECKING:
|
35
|
+
from collections.abc import Generator
|
36
|
+
|
37
|
+
from sqlalchemy.engine import Connection, Engine
|
38
|
+
from sqlalchemy.engine.cursor import CursorResult
|
39
|
+
from sqlalchemy.engine.reflection import Inspector
|
40
|
+
from sqlalchemy.sql.base import Executable
|
41
|
+
from sqlalchemy.sql.elements import TextClause
|
42
|
+
from sqlalchemy.sql.type_api import TypeEngine
|
43
|
+
|
44
|
+
from airbyte_cdk.sql.shared.catalog_providers import CatalogProvider
|
45
|
+
|
46
|
+
|
47
|
+
class SQLRuntimeError(Exception):
|
48
|
+
"""Raised when an SQL operation fails."""
|
49
|
+
|
50
|
+
|
51
|
+
class SqlConfig(BaseModel, abc.ABC):
|
52
|
+
"""Common configuration for SQL connections."""
|
53
|
+
|
54
|
+
schema_name: str = Field(default="airbyte_raw")
|
55
|
+
"""The name of the schema to write to."""
|
56
|
+
|
57
|
+
table_prefix: str | None = ""
|
58
|
+
"""A prefix to add to created table names."""
|
59
|
+
|
60
|
+
@abc.abstractmethod
|
61
|
+
def get_sql_alchemy_url(self) -> SecretString:
|
62
|
+
"""Returns a SQL Alchemy URL."""
|
63
|
+
...
|
64
|
+
|
65
|
+
@abc.abstractmethod
|
66
|
+
def get_database_name(self) -> str:
|
67
|
+
"""Return the name of the database."""
|
68
|
+
...
|
69
|
+
|
70
|
+
@property
|
71
|
+
def config_hash(self) -> str | None:
|
72
|
+
"""Return a unique one-way hash of the configuration.
|
73
|
+
|
74
|
+
The generic implementation uses the SQL Alchemy URL, schema name, and table prefix. Some
|
75
|
+
inputs may be redundant with the SQL Alchemy URL, but this does not hurt the hash
|
76
|
+
uniqueness.
|
77
|
+
|
78
|
+
In most cases, subclasses do not need to override this method.
|
79
|
+
"""
|
80
|
+
return one_way_hash(
|
81
|
+
SecretString(
|
82
|
+
":".join(
|
83
|
+
[
|
84
|
+
str(self.get_sql_alchemy_url()),
|
85
|
+
self.schema_name or "",
|
86
|
+
self.table_prefix or "",
|
87
|
+
]
|
88
|
+
)
|
89
|
+
)
|
90
|
+
)
|
91
|
+
|
92
|
+
def get_sql_engine(self) -> Engine:
|
93
|
+
"""Return a new SQL engine to use."""
|
94
|
+
return create_engine(
|
95
|
+
url=self.get_sql_alchemy_url(),
|
96
|
+
echo=DEBUG_MODE,
|
97
|
+
execution_options={
|
98
|
+
"schema_translate_map": {None: self.schema_name},
|
99
|
+
},
|
100
|
+
future=True,
|
101
|
+
)
|
102
|
+
|
103
|
+
def get_vendor_client(self) -> object:
|
104
|
+
"""Return the vendor-specific client object.
|
105
|
+
|
106
|
+
This is used for vendor-specific operations.
|
107
|
+
|
108
|
+
Raises `NotImplementedError` if a custom vendor client is not defined.
|
109
|
+
"""
|
110
|
+
raise NotImplementedError(
|
111
|
+
f"The type '{type(self).__name__}' does not define a custom client."
|
112
|
+
)
|
113
|
+
|
114
|
+
|
115
|
+
class SqlProcessorBase(abc.ABC):
|
116
|
+
"""A base class to be used for SQL Caches."""
|
117
|
+
|
118
|
+
type_converter_class: type[SQLTypeConverter] = SQLTypeConverter
|
119
|
+
"""The type converter class to use for converting JSON schema types to SQL types."""
|
120
|
+
|
121
|
+
normalizer = LowerCaseNormalizer
|
122
|
+
"""The name normalizer to user for table and column name normalization."""
|
123
|
+
|
124
|
+
supports_merge_insert = False
|
125
|
+
"""True if the database supports the MERGE INTO syntax."""
|
126
|
+
|
127
|
+
def __init__(
|
128
|
+
self,
|
129
|
+
*,
|
130
|
+
sql_config: SqlConfig,
|
131
|
+
catalog_provider: CatalogProvider,
|
132
|
+
) -> None:
|
133
|
+
"""Create a new SQL processor."""
|
134
|
+
self._sql_config: SqlConfig = sql_config
|
135
|
+
self._catalog_provider: CatalogProvider | None = catalog_provider
|
136
|
+
|
137
|
+
self._pending_state_messages: dict[str, list[AirbyteStateMessage]] = defaultdict(list, {})
|
138
|
+
self._finalized_state_messages: dict[
|
139
|
+
str,
|
140
|
+
list[AirbyteStateMessage],
|
141
|
+
] = defaultdict(list, {})
|
142
|
+
|
143
|
+
self._setup()
|
144
|
+
self.type_converter = self.type_converter_class()
|
145
|
+
self._cached_table_definitions: dict[str, sqlalchemy.Table] = {}
|
146
|
+
|
147
|
+
self._known_schemas_list: list[str] = []
|
148
|
+
self._ensure_schema_exists()
|
149
|
+
|
150
|
+
@property
|
151
|
+
def catalog_provider(
|
152
|
+
self,
|
153
|
+
) -> CatalogProvider:
|
154
|
+
"""Return the catalog manager.
|
155
|
+
|
156
|
+
Subclasses should set this property to a valid catalog manager instance if one
|
157
|
+
is not explicitly passed to the constructor.
|
158
|
+
|
159
|
+
Raises:
|
160
|
+
AirbyteInternalError: If the catalog manager is not set.
|
161
|
+
"""
|
162
|
+
if not self._catalog_provider:
|
163
|
+
raise exc.AirbyteInternalError(
|
164
|
+
message="Catalog manager should exist but does not.",
|
165
|
+
)
|
166
|
+
|
167
|
+
return self._catalog_provider
|
168
|
+
|
169
|
+
def _setup(self) -> None: # noqa: B027 # Intentionally empty, not abstract
|
170
|
+
"""Create the database.
|
171
|
+
|
172
|
+
By default this is a no-op but subclasses can override this method to prepare
|
173
|
+
any necessary resources.
|
174
|
+
"""
|
175
|
+
pass
|
176
|
+
|
177
|
+
def _do_checkpoint( # noqa: B027 # Intentionally empty, not abstract
|
178
|
+
self,
|
179
|
+
connection: Connection | None = None,
|
180
|
+
) -> None:
|
181
|
+
"""Checkpoint the given connection.
|
182
|
+
|
183
|
+
If the WAL log needs to be, it will be flushed.
|
184
|
+
|
185
|
+
For most SQL databases, this is a no-op. However, it exists so that
|
186
|
+
subclasses can override this method to perform a checkpoint operation.
|
187
|
+
"""
|
188
|
+
pass
|
189
|
+
|
190
|
+
# Public interface:
|
191
|
+
|
192
|
+
@property
|
193
|
+
def sql_config(self) -> SqlConfig:
|
194
|
+
"""Return the SQL configuration."""
|
195
|
+
return self._sql_config
|
196
|
+
|
197
|
+
def get_sql_alchemy_url(self) -> SecretString:
|
198
|
+
"""Return the SQLAlchemy URL to use."""
|
199
|
+
return self.sql_config.get_sql_alchemy_url()
|
200
|
+
|
201
|
+
@final
|
202
|
+
@cached_property
|
203
|
+
def database_name(self) -> str:
|
204
|
+
"""Return the name of the database."""
|
205
|
+
return self.sql_config.get_database_name()
|
206
|
+
|
207
|
+
@final
|
208
|
+
def get_sql_engine(self) -> Engine:
|
209
|
+
"""Return a new SQL engine to use."""
|
210
|
+
return self.sql_config.get_sql_engine()
|
211
|
+
|
212
|
+
@contextmanager
|
213
|
+
def get_sql_connection(self) -> Generator[sqlalchemy.engine.Connection, None, None]:
|
214
|
+
"""A context manager which returns a new SQL connection for running queries.
|
215
|
+
|
216
|
+
If the connection needs to close, it will be closed automatically.
|
217
|
+
"""
|
218
|
+
with self.get_sql_engine().begin() as connection:
|
219
|
+
self._init_connection_settings(connection)
|
220
|
+
yield connection
|
221
|
+
|
222
|
+
connection.close()
|
223
|
+
del connection
|
224
|
+
|
225
|
+
def get_sql_table_name(
|
226
|
+
self,
|
227
|
+
stream_name: str,
|
228
|
+
) -> str:
|
229
|
+
"""Return the name of the SQL table for the given stream."""
|
230
|
+
table_prefix = self.sql_config.table_prefix
|
231
|
+
return self.normalizer.normalize(
|
232
|
+
f"{table_prefix}{stream_name}",
|
233
|
+
)
|
234
|
+
|
235
|
+
@final
|
236
|
+
def get_sql_table(
|
237
|
+
self,
|
238
|
+
stream_name: str,
|
239
|
+
) -> sqlalchemy.Table:
|
240
|
+
"""Return the main table object for the stream."""
|
241
|
+
return self._get_table_by_name(
|
242
|
+
self.get_sql_table_name(stream_name),
|
243
|
+
)
|
244
|
+
|
245
|
+
# Protected members (non-public interface):
|
246
|
+
|
247
|
+
def _init_connection_settings(self, connection: Connection) -> None: # noqa: B027 # Intentionally empty, not abstract
|
248
|
+
"""This is called automatically whenever a new connection is created.
|
249
|
+
|
250
|
+
By default this is a no-op. Subclasses can use this to set connection settings, such as
|
251
|
+
timezone, case-sensitivity settings, and other session-level variables.
|
252
|
+
"""
|
253
|
+
pass
|
254
|
+
|
255
|
+
def _invalidate_table_cache(
|
256
|
+
self,
|
257
|
+
table_name: str,
|
258
|
+
) -> None:
|
259
|
+
"""Invalidate the the named table cache.
|
260
|
+
|
261
|
+
This should be called whenever the table schema is known to have changed.
|
262
|
+
"""
|
263
|
+
if table_name in self._cached_table_definitions:
|
264
|
+
del self._cached_table_definitions[table_name]
|
265
|
+
|
266
|
+
def _get_table_by_name(
|
267
|
+
self,
|
268
|
+
table_name: str,
|
269
|
+
*,
|
270
|
+
force_refresh: bool = False,
|
271
|
+
shallow_okay: bool = False,
|
272
|
+
) -> sqlalchemy.Table:
|
273
|
+
"""Return a table object from a table name.
|
274
|
+
|
275
|
+
If 'shallow_okay' is True, the table will be returned without requiring properties to
|
276
|
+
be read from the database.
|
277
|
+
|
278
|
+
To prevent unnecessary round-trips to the database, the table is cached after the first
|
279
|
+
query. To ignore the cache and force a refresh, set 'force_refresh' to True.
|
280
|
+
"""
|
281
|
+
if force_refresh and shallow_okay:
|
282
|
+
raise exc.AirbyteInternalError(
|
283
|
+
message="Cannot force refresh and use shallow query at the same time."
|
284
|
+
)
|
285
|
+
|
286
|
+
if force_refresh and table_name in self._cached_table_definitions:
|
287
|
+
self._invalidate_table_cache(table_name)
|
288
|
+
|
289
|
+
if table_name not in self._cached_table_definitions:
|
290
|
+
if shallow_okay:
|
291
|
+
# Return a shallow instance, without column declarations. Do not cache
|
292
|
+
# the table definition in this case.
|
293
|
+
return sqlalchemy.Table(
|
294
|
+
table_name,
|
295
|
+
sqlalchemy.MetaData(schema=self.sql_config.schema_name),
|
296
|
+
)
|
297
|
+
|
298
|
+
self._cached_table_definitions[table_name] = sqlalchemy.Table(
|
299
|
+
table_name,
|
300
|
+
sqlalchemy.MetaData(schema=self.sql_config.schema_name),
|
301
|
+
autoload_with=self.get_sql_engine(),
|
302
|
+
)
|
303
|
+
|
304
|
+
return self._cached_table_definitions[table_name]
|
305
|
+
|
306
|
+
def _ensure_schema_exists(
|
307
|
+
self,
|
308
|
+
) -> None:
|
309
|
+
schema_name = self.normalizer.normalize(self.sql_config.schema_name)
|
310
|
+
known_schemas_list = self.normalizer.normalize_list(self._known_schemas_list)
|
311
|
+
if known_schemas_list and schema_name in known_schemas_list:
|
312
|
+
return # Already exists
|
313
|
+
|
314
|
+
schemas_list = self.normalizer.normalize_list(self._get_schemas_list())
|
315
|
+
if schema_name in schemas_list:
|
316
|
+
return
|
317
|
+
|
318
|
+
sql = f"CREATE SCHEMA IF NOT EXISTS {schema_name}"
|
319
|
+
|
320
|
+
try:
|
321
|
+
self._execute_sql(sql)
|
322
|
+
except Exception as ex:
|
323
|
+
# Ignore schema exists errors.
|
324
|
+
if "already exists" not in str(ex):
|
325
|
+
raise
|
326
|
+
|
327
|
+
if DEBUG_MODE:
|
328
|
+
found_schemas = schemas_list
|
329
|
+
assert (
|
330
|
+
schema_name in found_schemas
|
331
|
+
), f"Schema {schema_name} was not created. Found: {found_schemas}"
|
332
|
+
|
333
|
+
def _quote_identifier(self, identifier: str) -> str:
|
334
|
+
"""Return the given identifier, quoted."""
|
335
|
+
return f'"{identifier}"'
|
336
|
+
|
337
|
+
@final
|
338
|
+
def _get_temp_table_name(
|
339
|
+
self,
|
340
|
+
stream_name: str,
|
341
|
+
batch_id: str | None = None, # ULID of the batch
|
342
|
+
) -> str:
|
343
|
+
"""Return a new (unique) temporary table name."""
|
344
|
+
if not batch_id:
|
345
|
+
batch_id = str(ulid.ULID())
|
346
|
+
|
347
|
+
# Use the first 6 and last 3 characters of the ULID. This gives great uniqueness while
|
348
|
+
# limiting the table name suffix to 10 characters, including the underscore.
|
349
|
+
suffix = f"{batch_id[:6]}{batch_id[-3:]}" if len(batch_id) > 9 else batch_id # noqa: PLR2004 # Allow magic int value
|
350
|
+
|
351
|
+
# Note: The normalizer may truncate the table name if the database has a name length limit.
|
352
|
+
# For instance, the Postgres normalizer will enforce a 63-character limit on table names.
|
353
|
+
return self.normalizer.normalize(f"{stream_name}_{suffix}")
|
354
|
+
|
355
|
+
def _fully_qualified(
|
356
|
+
self,
|
357
|
+
table_name: str,
|
358
|
+
) -> str:
|
359
|
+
"""Return the fully qualified name of the given table."""
|
360
|
+
return f"{self.sql_config.schema_name}.{self._quote_identifier(table_name)}"
|
361
|
+
|
362
|
+
@final
|
363
|
+
def _create_table_for_loading(
|
364
|
+
self,
|
365
|
+
/,
|
366
|
+
stream_name: str,
|
367
|
+
batch_id: str | None,
|
368
|
+
) -> str:
|
369
|
+
"""Create a new table for loading data."""
|
370
|
+
temp_table_name = self._get_temp_table_name(stream_name, batch_id)
|
371
|
+
column_definition_str = ",\n ".join(
|
372
|
+
f"{self._quote_identifier(column_name)} {sql_type}"
|
373
|
+
for column_name, sql_type in self._get_sql_column_definitions(stream_name).items()
|
374
|
+
)
|
375
|
+
self._create_table(temp_table_name, column_definition_str)
|
376
|
+
|
377
|
+
return temp_table_name
|
378
|
+
|
379
|
+
def _get_tables_list(
|
380
|
+
self,
|
381
|
+
) -> list[str]:
|
382
|
+
"""Return a list of all tables in the database."""
|
383
|
+
with self.get_sql_connection() as conn:
|
384
|
+
inspector: Inspector = sqlalchemy.inspect(conn)
|
385
|
+
return inspector.get_table_names(schema=self.sql_config.schema_name) # type: ignore
|
386
|
+
|
387
|
+
def _get_schemas_list(
|
388
|
+
self,
|
389
|
+
database_name: str | None = None,
|
390
|
+
*,
|
391
|
+
force_refresh: bool = False,
|
392
|
+
) -> list[str]:
|
393
|
+
"""Return a list of all tables in the database."""
|
394
|
+
if not force_refresh and self._known_schemas_list:
|
395
|
+
return self._known_schemas_list
|
396
|
+
|
397
|
+
inspector: Inspector = sqlalchemy.inspect(self.get_sql_engine())
|
398
|
+
database_name = database_name or self.database_name
|
399
|
+
found_schemas = inspector.get_schema_names()
|
400
|
+
self._known_schemas_list = [
|
401
|
+
found_schema.split(".")[-1].strip('"')
|
402
|
+
for found_schema in found_schemas
|
403
|
+
if "." not in found_schema
|
404
|
+
or (found_schema.split(".")[0].lower().strip('"') == database_name.lower())
|
405
|
+
]
|
406
|
+
return self._known_schemas_list
|
407
|
+
|
408
|
+
def _ensure_final_table_exists(
|
409
|
+
self,
|
410
|
+
stream_name: str,
|
411
|
+
*,
|
412
|
+
create_if_missing: bool = True,
|
413
|
+
) -> str:
|
414
|
+
"""Create the final table if it doesn't already exist.
|
415
|
+
|
416
|
+
Return the table name.
|
417
|
+
"""
|
418
|
+
table_name = self.get_sql_table_name(stream_name)
|
419
|
+
did_exist = self._table_exists(table_name)
|
420
|
+
if not did_exist and create_if_missing:
|
421
|
+
column_definition_str = ",\n ".join(
|
422
|
+
f"{self._quote_identifier(column_name)} {sql_type}"
|
423
|
+
for column_name, sql_type in self._get_sql_column_definitions(
|
424
|
+
stream_name,
|
425
|
+
).items()
|
426
|
+
)
|
427
|
+
self._create_table(table_name, column_definition_str)
|
428
|
+
|
429
|
+
return table_name
|
430
|
+
|
431
|
+
def _ensure_compatible_table_schema(
|
432
|
+
self,
|
433
|
+
stream_name: str,
|
434
|
+
table_name: str,
|
435
|
+
) -> None:
|
436
|
+
"""Return true if the given table is compatible with the stream's schema.
|
437
|
+
|
438
|
+
Raises an exception if the table schema is not compatible with the schema of the
|
439
|
+
input stream.
|
440
|
+
"""
|
441
|
+
# TODO: Expand this to check for column types and sizes.
|
442
|
+
# https://github.com/airbytehq/Airbyte/issues/321
|
443
|
+
self._add_missing_columns_to_table(
|
444
|
+
stream_name=stream_name,
|
445
|
+
table_name=table_name,
|
446
|
+
)
|
447
|
+
|
448
|
+
@final
|
449
|
+
def _create_table(
|
450
|
+
self,
|
451
|
+
table_name: str,
|
452
|
+
column_definition_str: str,
|
453
|
+
primary_keys: list[str] | None = None,
|
454
|
+
) -> None:
|
455
|
+
if primary_keys:
|
456
|
+
pk_str = ", ".join(primary_keys)
|
457
|
+
column_definition_str += f",\n PRIMARY KEY ({pk_str})"
|
458
|
+
|
459
|
+
cmd = f"""
|
460
|
+
CREATE TABLE {self._fully_qualified(table_name)} (
|
461
|
+
{column_definition_str}
|
462
|
+
)
|
463
|
+
"""
|
464
|
+
_ = self._execute_sql(cmd)
|
465
|
+
|
466
|
+
@final
|
467
|
+
def _get_sql_column_definitions(
|
468
|
+
self,
|
469
|
+
stream_name: str,
|
470
|
+
) -> dict[str, sqlalchemy.types.TypeEngine[Any]]:
|
471
|
+
"""Return the column definitions for the given stream."""
|
472
|
+
columns: dict[str, sqlalchemy.types.TypeEngine[Any]] = {}
|
473
|
+
properties = self.catalog_provider.get_stream_properties(stream_name)
|
474
|
+
for property_name, json_schema_property_def in properties.items():
|
475
|
+
clean_prop_name = self.normalizer.normalize(property_name)
|
476
|
+
columns[clean_prop_name] = self.type_converter.to_sql_type(
|
477
|
+
json_schema_property_def,
|
478
|
+
)
|
479
|
+
|
480
|
+
columns[AB_RAW_ID_COLUMN] = self.type_converter_class.get_string_type()
|
481
|
+
columns[AB_EXTRACTED_AT_COLUMN] = sqlalchemy.TIMESTAMP()
|
482
|
+
columns[AB_META_COLUMN] = self.type_converter_class.get_json_type()
|
483
|
+
|
484
|
+
return columns
|
485
|
+
|
486
|
+
def _execute_sql(self, sql: str | TextClause | Executable) -> CursorResult[Any]:
|
487
|
+
"""Execute the given SQL statement."""
|
488
|
+
if isinstance(sql, str):
|
489
|
+
sql = text(sql)
|
490
|
+
|
491
|
+
with self.get_sql_connection() as conn:
|
492
|
+
try:
|
493
|
+
result = conn.execute(sql)
|
494
|
+
except (
|
495
|
+
ProgrammingError,
|
496
|
+
SQLAlchemyError,
|
497
|
+
) as ex:
|
498
|
+
msg = f"Error when executing SQL:\n{sql}\n{type(ex).__name__}{ex!s}"
|
499
|
+
raise SQLRuntimeError(msg) from None # from ex
|
500
|
+
|
501
|
+
return result
|
502
|
+
|
503
|
+
def _drop_temp_table(
|
504
|
+
self,
|
505
|
+
table_name: str,
|
506
|
+
*,
|
507
|
+
if_exists: bool = True,
|
508
|
+
) -> None:
|
509
|
+
"""Drop the given table."""
|
510
|
+
exists_str = "IF EXISTS" if if_exists else ""
|
511
|
+
self._execute_sql(f"DROP TABLE {exists_str} {self._fully_qualified(table_name)}")
|
512
|
+
|
513
|
+
def _write_files_to_new_table(
|
514
|
+
self,
|
515
|
+
files: list[Path],
|
516
|
+
stream_name: str,
|
517
|
+
batch_id: str,
|
518
|
+
) -> str:
|
519
|
+
"""Write a file(s) to a new table.
|
520
|
+
|
521
|
+
This is a generic implementation, which can be overridden by subclasses
|
522
|
+
to improve performance.
|
523
|
+
"""
|
524
|
+
temp_table_name = self._create_table_for_loading(stream_name, batch_id)
|
525
|
+
for file_path in files:
|
526
|
+
dataframe = pd.read_json(file_path, lines=True)
|
527
|
+
|
528
|
+
sql_column_definitions: dict[str, TypeEngine[Any]] = self._get_sql_column_definitions(
|
529
|
+
stream_name
|
530
|
+
)
|
531
|
+
|
532
|
+
# Remove fields that are not in the schema
|
533
|
+
for col_name in dataframe.columns:
|
534
|
+
if col_name not in sql_column_definitions:
|
535
|
+
dataframe = dataframe.drop(columns=col_name)
|
536
|
+
|
537
|
+
# Pandas will auto-create the table if it doesn't exist, which we don't want.
|
538
|
+
if not self._table_exists(temp_table_name):
|
539
|
+
raise exc.AirbyteInternalError(
|
540
|
+
message="Table does not exist after creation.",
|
541
|
+
context={
|
542
|
+
"temp_table_name": temp_table_name,
|
543
|
+
},
|
544
|
+
)
|
545
|
+
|
546
|
+
# Normalize all column names to lower case.
|
547
|
+
dataframe.columns = Index([self.normalizer.normalize(col) for col in dataframe.columns])
|
548
|
+
|
549
|
+
# Write the data to the table.
|
550
|
+
dataframe.to_sql(
|
551
|
+
temp_table_name,
|
552
|
+
self.get_sql_alchemy_url(),
|
553
|
+
schema=self.sql_config.schema_name,
|
554
|
+
if_exists="append",
|
555
|
+
index=False,
|
556
|
+
dtype=sql_column_definitions, # type: ignore[arg-type]
|
557
|
+
)
|
558
|
+
return temp_table_name
|
559
|
+
|
560
|
+
def _add_column_to_table(
|
561
|
+
self,
|
562
|
+
table: Table,
|
563
|
+
column_name: str,
|
564
|
+
column_type: sqlalchemy.types.TypeEngine[Any],
|
565
|
+
) -> None:
|
566
|
+
"""Add a column to the given table."""
|
567
|
+
self._execute_sql(
|
568
|
+
text(
|
569
|
+
f"ALTER TABLE {self._fully_qualified(table.name)} "
|
570
|
+
f"ADD COLUMN {column_name} {column_type}"
|
571
|
+
),
|
572
|
+
)
|
573
|
+
|
574
|
+
def _add_missing_columns_to_table(
|
575
|
+
self,
|
576
|
+
stream_name: str,
|
577
|
+
table_name: str,
|
578
|
+
) -> None:
|
579
|
+
"""Add missing columns to the table.
|
580
|
+
|
581
|
+
This is a no-op if all columns are already present.
|
582
|
+
"""
|
583
|
+
columns = self._get_sql_column_definitions(stream_name)
|
584
|
+
# First check without forcing a refresh of the cache (faster). If nothing is missing,
|
585
|
+
# then we're done.
|
586
|
+
table = self._get_table_by_name(
|
587
|
+
table_name,
|
588
|
+
force_refresh=False,
|
589
|
+
)
|
590
|
+
missing_columns: bool = any(column_name not in table.columns for column_name in columns)
|
591
|
+
|
592
|
+
if missing_columns:
|
593
|
+
# If we found missing columns, refresh the cache and then take action on anything
|
594
|
+
# that's still confirmed missing.
|
595
|
+
columns_added = False
|
596
|
+
table = self._get_table_by_name(
|
597
|
+
table_name,
|
598
|
+
force_refresh=True,
|
599
|
+
)
|
600
|
+
for column_name, column_type in columns.items():
|
601
|
+
if column_name not in table.columns:
|
602
|
+
self._add_column_to_table(table, column_name, column_type)
|
603
|
+
columns_added = True
|
604
|
+
|
605
|
+
if columns_added:
|
606
|
+
# We've added columns, so invalidate the cache.
|
607
|
+
self._invalidate_table_cache(table_name)
|
608
|
+
|
609
|
+
def _append_temp_table_to_final_table(
|
610
|
+
self,
|
611
|
+
temp_table_name: str,
|
612
|
+
final_table_name: str,
|
613
|
+
stream_name: str,
|
614
|
+
) -> None:
|
615
|
+
nl = "\n"
|
616
|
+
columns = [self._quote_identifier(c) for c in self._get_sql_column_definitions(stream_name)]
|
617
|
+
self._execute_sql(
|
618
|
+
f"""
|
619
|
+
INSERT INTO {self._fully_qualified(final_table_name)} (
|
620
|
+
{f',{nl} '.join(columns)}
|
621
|
+
)
|
622
|
+
SELECT
|
623
|
+
{f',{nl} '.join(columns)}
|
624
|
+
FROM {self._fully_qualified(temp_table_name)}
|
625
|
+
""",
|
626
|
+
)
|
627
|
+
|
628
|
+
def _swap_temp_table_with_final_table(
|
629
|
+
self,
|
630
|
+
stream_name: str,
|
631
|
+
temp_table_name: str,
|
632
|
+
final_table_name: str,
|
633
|
+
) -> None:
|
634
|
+
"""Merge the temp table into the main one.
|
635
|
+
|
636
|
+
This implementation requires MERGE support in the SQL DB.
|
637
|
+
Databases that do not support this syntax can override this method.
|
638
|
+
"""
|
639
|
+
if final_table_name is None:
|
640
|
+
raise exc.AirbyteInternalError(message="Arg 'final_table_name' cannot be None.")
|
641
|
+
if temp_table_name is None:
|
642
|
+
raise exc.AirbyteInternalError(message="Arg 'temp_table_name' cannot be None.")
|
643
|
+
|
644
|
+
_ = stream_name
|
645
|
+
deletion_name = f"{final_table_name}_deleteme"
|
646
|
+
commands = "\n".join(
|
647
|
+
[
|
648
|
+
f"ALTER TABLE {self._fully_qualified(final_table_name)} RENAME "
|
649
|
+
f"TO {deletion_name};",
|
650
|
+
f"ALTER TABLE {self._fully_qualified(temp_table_name)} RENAME "
|
651
|
+
f"TO {final_table_name};",
|
652
|
+
f"DROP TABLE {self._fully_qualified(deletion_name)};",
|
653
|
+
]
|
654
|
+
)
|
655
|
+
self._execute_sql(commands)
|
656
|
+
|
657
|
+
def _merge_temp_table_to_final_table(
|
658
|
+
self,
|
659
|
+
stream_name: str,
|
660
|
+
temp_table_name: str,
|
661
|
+
final_table_name: str,
|
662
|
+
) -> None:
|
663
|
+
"""Merge the temp table into the main one.
|
664
|
+
|
665
|
+
This implementation requires MERGE support in the SQL DB.
|
666
|
+
Databases that do not support this syntax can override this method.
|
667
|
+
"""
|
668
|
+
nl = "\n"
|
669
|
+
columns = {self._quote_identifier(c) for c in self._get_sql_column_definitions(stream_name)}
|
670
|
+
pk_columns = {
|
671
|
+
self._quote_identifier(c) for c in self.catalog_provider.get_primary_keys(stream_name)
|
672
|
+
}
|
673
|
+
non_pk_columns = columns - pk_columns
|
674
|
+
join_clause = f"{nl} AND ".join(f"tmp.{pk_col} = final.{pk_col}" for pk_col in pk_columns)
|
675
|
+
set_clause = f"{nl} , ".join(f"{col} = tmp.{col}" for col in non_pk_columns)
|
676
|
+
self._execute_sql(
|
677
|
+
f"""
|
678
|
+
MERGE INTO {self._fully_qualified(final_table_name)} final
|
679
|
+
USING (
|
680
|
+
SELECT *
|
681
|
+
FROM {self._fully_qualified(temp_table_name)}
|
682
|
+
) AS tmp
|
683
|
+
ON {join_clause}
|
684
|
+
WHEN MATCHED THEN UPDATE
|
685
|
+
SET
|
686
|
+
{set_clause}
|
687
|
+
WHEN NOT MATCHED THEN INSERT
|
688
|
+
(
|
689
|
+
{f',{nl} '.join(columns)}
|
690
|
+
)
|
691
|
+
VALUES (
|
692
|
+
tmp.{f',{nl} tmp.'.join(columns)}
|
693
|
+
);
|
694
|
+
""",
|
695
|
+
)
|
696
|
+
|
697
|
+
def _get_column_by_name(self, table: str | Table, column_name: str) -> Column[Any]:
|
698
|
+
"""Return the column object for the given column name.
|
699
|
+
|
700
|
+
This method is case-insensitive.
|
701
|
+
"""
|
702
|
+
if isinstance(table, str):
|
703
|
+
table = self._get_table_by_name(table)
|
704
|
+
try:
|
705
|
+
# Try to get the column in a case-insensitive manner
|
706
|
+
return next(col for col in table.c if col.name.lower() == column_name.lower())
|
707
|
+
except StopIteration:
|
708
|
+
raise exc.AirbyteInternalError(
|
709
|
+
message="Could not find matching column.",
|
710
|
+
context={
|
711
|
+
"table": table,
|
712
|
+
"column_name": column_name,
|
713
|
+
},
|
714
|
+
) from None
|
715
|
+
|
716
|
+
def _emulated_merge_temp_table_to_final_table(
|
717
|
+
self,
|
718
|
+
stream_name: str,
|
719
|
+
temp_table_name: str,
|
720
|
+
final_table_name: str,
|
721
|
+
) -> None:
|
722
|
+
"""Emulate the merge operation using a series of SQL commands.
|
723
|
+
|
724
|
+
This is a fallback implementation for databases that do not support MERGE.
|
725
|
+
"""
|
726
|
+
final_table = self._get_table_by_name(final_table_name)
|
727
|
+
temp_table = self._get_table_by_name(temp_table_name)
|
728
|
+
pk_columns = self.catalog_provider.get_primary_keys(stream_name)
|
729
|
+
|
730
|
+
columns_to_update: set[str] = self._get_sql_column_definitions(
|
731
|
+
stream_name=stream_name
|
732
|
+
).keys() - set(pk_columns)
|
733
|
+
|
734
|
+
# Create a dictionary mapping columns in users_final to users_stage for updating
|
735
|
+
update_values = {
|
736
|
+
self._get_column_by_name(final_table, column): (
|
737
|
+
self._get_column_by_name(temp_table, column)
|
738
|
+
)
|
739
|
+
for column in columns_to_update
|
740
|
+
}
|
741
|
+
|
742
|
+
# Craft the WHERE clause for composite primary keys
|
743
|
+
join_conditions = [
|
744
|
+
self._get_column_by_name(final_table, pk_column)
|
745
|
+
== self._get_column_by_name(temp_table, pk_column)
|
746
|
+
for pk_column in pk_columns
|
747
|
+
]
|
748
|
+
join_clause = and_(*join_conditions)
|
749
|
+
|
750
|
+
# Craft the UPDATE statement
|
751
|
+
update_stmt = update(final_table).values(update_values).where(join_clause)
|
752
|
+
|
753
|
+
# Define a join between temp_table and final_table
|
754
|
+
joined_table = temp_table.outerjoin(final_table, join_clause)
|
755
|
+
|
756
|
+
# Define a condition that checks for records in temp_table that do not have a corresponding
|
757
|
+
# record in final_table
|
758
|
+
where_not_exists_clause = self._get_column_by_name(final_table, pk_columns[0]) == null()
|
759
|
+
|
760
|
+
# Select records from temp_table that are not in final_table
|
761
|
+
select_new_records_stmt = (
|
762
|
+
select(temp_table).select_from(joined_table).where(where_not_exists_clause)
|
763
|
+
)
|
764
|
+
|
765
|
+
# Craft the INSERT statement using the select statement
|
766
|
+
insert_new_records_stmt = insert(final_table).from_select(
|
767
|
+
names=[column.name for column in temp_table.columns], select=select_new_records_stmt
|
768
|
+
)
|
769
|
+
|
770
|
+
if DEBUG_MODE:
|
771
|
+
print(str(update_stmt))
|
772
|
+
print(str(insert_new_records_stmt))
|
773
|
+
|
774
|
+
with self.get_sql_connection() as conn:
|
775
|
+
conn.execute(update_stmt)
|
776
|
+
conn.execute(insert_new_records_stmt)
|
777
|
+
|
778
|
+
def _table_exists(
|
779
|
+
self,
|
780
|
+
table_name: str,
|
781
|
+
) -> bool:
|
782
|
+
"""Return true if the given table exists.
|
783
|
+
|
784
|
+
Subclasses may override this method to provide a more efficient implementation.
|
785
|
+
"""
|
786
|
+
return table_name in self._get_tables_list()
|