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.
Files changed (368) hide show
  1. airbyte_cdk/__init__.py +358 -0
  2. airbyte_cdk/cli/__init__.py +1 -0
  3. airbyte_cdk/cli/source_declarative_manifest/__init__.py +5 -0
  4. airbyte_cdk/cli/source_declarative_manifest/_run.py +236 -0
  5. airbyte_cdk/cli/source_declarative_manifest/spec.json +17 -0
  6. airbyte_cdk/config_observation.py +104 -0
  7. airbyte_cdk/connector.py +123 -0
  8. airbyte_cdk/connector_builder/README.md +53 -0
  9. airbyte_cdk/connector_builder/__init__.py +3 -0
  10. airbyte_cdk/connector_builder/connector_builder_handler.py +121 -0
  11. airbyte_cdk/connector_builder/main.py +107 -0
  12. airbyte_cdk/connector_builder/models.py +73 -0
  13. airbyte_cdk/connector_builder/test_reader/__init__.py +7 -0
  14. airbyte_cdk/connector_builder/test_reader/helpers.py +689 -0
  15. airbyte_cdk/connector_builder/test_reader/message_grouper.py +173 -0
  16. airbyte_cdk/connector_builder/test_reader/reader.py +441 -0
  17. airbyte_cdk/connector_builder/test_reader/types.py +83 -0
  18. airbyte_cdk/destinations/__init__.py +8 -0
  19. airbyte_cdk/destinations/destination.py +154 -0
  20. airbyte_cdk/destinations/vector_db_based/README.md +37 -0
  21. airbyte_cdk/destinations/vector_db_based/__init__.py +38 -0
  22. airbyte_cdk/destinations/vector_db_based/config.py +298 -0
  23. airbyte_cdk/destinations/vector_db_based/document_processor.py +223 -0
  24. airbyte_cdk/destinations/vector_db_based/embedder.py +303 -0
  25. airbyte_cdk/destinations/vector_db_based/indexer.py +78 -0
  26. airbyte_cdk/destinations/vector_db_based/test_utils.py +63 -0
  27. airbyte_cdk/destinations/vector_db_based/utils.py +35 -0
  28. airbyte_cdk/destinations/vector_db_based/writer.py +104 -0
  29. airbyte_cdk/entrypoint.py +414 -0
  30. airbyte_cdk/exception_handler.py +56 -0
  31. airbyte_cdk/logger.py +109 -0
  32. airbyte_cdk/models/__init__.py +72 -0
  33. airbyte_cdk/models/airbyte_protocol.py +88 -0
  34. airbyte_cdk/models/airbyte_protocol_serializers.py +44 -0
  35. airbyte_cdk/models/well_known_types.py +5 -0
  36. airbyte_cdk/py.typed +0 -0
  37. airbyte_cdk/sources/__init__.py +26 -0
  38. airbyte_cdk/sources/abstract_source.py +326 -0
  39. airbyte_cdk/sources/concurrent_source/__init__.py +8 -0
  40. airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py +255 -0
  41. airbyte_cdk/sources/concurrent_source/concurrent_source.py +165 -0
  42. airbyte_cdk/sources/concurrent_source/concurrent_source_adapter.py +147 -0
  43. airbyte_cdk/sources/concurrent_source/partition_generation_completed_sentinel.py +24 -0
  44. airbyte_cdk/sources/concurrent_source/stream_thread_exception.py +25 -0
  45. airbyte_cdk/sources/concurrent_source/thread_pool_manager.py +115 -0
  46. airbyte_cdk/sources/config.py +27 -0
  47. airbyte_cdk/sources/connector_state_manager.py +161 -0
  48. airbyte_cdk/sources/declarative/__init__.py +3 -0
  49. airbyte_cdk/sources/declarative/async_job/__init__.py +0 -0
  50. airbyte_cdk/sources/declarative/async_job/job.py +52 -0
  51. airbyte_cdk/sources/declarative/async_job/job_orchestrator.py +525 -0
  52. airbyte_cdk/sources/declarative/async_job/job_tracker.py +79 -0
  53. airbyte_cdk/sources/declarative/async_job/repository.py +35 -0
  54. airbyte_cdk/sources/declarative/async_job/status.py +24 -0
  55. airbyte_cdk/sources/declarative/async_job/timer.py +39 -0
  56. airbyte_cdk/sources/declarative/auth/__init__.py +8 -0
  57. airbyte_cdk/sources/declarative/auth/declarative_authenticator.py +42 -0
  58. airbyte_cdk/sources/declarative/auth/jwt.py +197 -0
  59. airbyte_cdk/sources/declarative/auth/oauth.py +293 -0
  60. airbyte_cdk/sources/declarative/auth/selective_authenticator.py +45 -0
  61. airbyte_cdk/sources/declarative/auth/token.py +267 -0
  62. airbyte_cdk/sources/declarative/auth/token_provider.py +82 -0
  63. airbyte_cdk/sources/declarative/checks/__init__.py +24 -0
  64. airbyte_cdk/sources/declarative/checks/check_dynamic_stream.py +61 -0
  65. airbyte_cdk/sources/declarative/checks/check_stream.py +56 -0
  66. airbyte_cdk/sources/declarative/checks/connection_checker.py +35 -0
  67. airbyte_cdk/sources/declarative/concurrency_level/__init__.py +7 -0
  68. airbyte_cdk/sources/declarative/concurrency_level/concurrency_level.py +50 -0
  69. airbyte_cdk/sources/declarative/concurrent_declarative_source.py +526 -0
  70. airbyte_cdk/sources/declarative/datetime/__init__.py +3 -0
  71. airbyte_cdk/sources/declarative/datetime/datetime_parser.py +65 -0
  72. airbyte_cdk/sources/declarative/datetime/min_max_datetime.py +118 -0
  73. airbyte_cdk/sources/declarative/declarative_component_schema.yaml +3975 -0
  74. airbyte_cdk/sources/declarative/declarative_source.py +36 -0
  75. airbyte_cdk/sources/declarative/declarative_stream.py +241 -0
  76. airbyte_cdk/sources/declarative/decoders/__init__.py +33 -0
  77. airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py +218 -0
  78. airbyte_cdk/sources/declarative/decoders/decoder.py +32 -0
  79. airbyte_cdk/sources/declarative/decoders/decoder_parser.py +30 -0
  80. airbyte_cdk/sources/declarative/decoders/json_decoder.py +65 -0
  81. airbyte_cdk/sources/declarative/decoders/noop_decoder.py +21 -0
  82. airbyte_cdk/sources/declarative/decoders/pagination_decoder_decorator.py +39 -0
  83. airbyte_cdk/sources/declarative/decoders/xml_decoder.py +98 -0
  84. airbyte_cdk/sources/declarative/decoders/zipfile_decoder.py +56 -0
  85. airbyte_cdk/sources/declarative/exceptions.py +9 -0
  86. airbyte_cdk/sources/declarative/extractors/__init__.py +21 -0
  87. airbyte_cdk/sources/declarative/extractors/dpath_extractor.py +86 -0
  88. airbyte_cdk/sources/declarative/extractors/http_selector.py +37 -0
  89. airbyte_cdk/sources/declarative/extractors/record_extractor.py +27 -0
  90. airbyte_cdk/sources/declarative/extractors/record_filter.py +91 -0
  91. airbyte_cdk/sources/declarative/extractors/record_selector.py +170 -0
  92. airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py +176 -0
  93. airbyte_cdk/sources/declarative/extractors/type_transformer.py +55 -0
  94. airbyte_cdk/sources/declarative/incremental/__init__.py +37 -0
  95. airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py +497 -0
  96. airbyte_cdk/sources/declarative/incremental/datetime_based_cursor.py +459 -0
  97. airbyte_cdk/sources/declarative/incremental/declarative_cursor.py +13 -0
  98. airbyte_cdk/sources/declarative/incremental/global_substream_cursor.py +357 -0
  99. airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py +380 -0
  100. airbyte_cdk/sources/declarative/incremental/per_partition_with_global.py +200 -0
  101. airbyte_cdk/sources/declarative/incremental/resumable_full_refresh_cursor.py +122 -0
  102. airbyte_cdk/sources/declarative/interpolation/__init__.py +9 -0
  103. airbyte_cdk/sources/declarative/interpolation/filters.py +139 -0
  104. airbyte_cdk/sources/declarative/interpolation/interpolated_boolean.py +66 -0
  105. airbyte_cdk/sources/declarative/interpolation/interpolated_mapping.py +56 -0
  106. airbyte_cdk/sources/declarative/interpolation/interpolated_nested_mapping.py +52 -0
  107. airbyte_cdk/sources/declarative/interpolation/interpolated_string.py +79 -0
  108. airbyte_cdk/sources/declarative/interpolation/interpolation.py +34 -0
  109. airbyte_cdk/sources/declarative/interpolation/jinja.py +161 -0
  110. airbyte_cdk/sources/declarative/interpolation/macros.py +191 -0
  111. airbyte_cdk/sources/declarative/manifest_declarative_source.py +421 -0
  112. airbyte_cdk/sources/declarative/migrations/__init__.py +0 -0
  113. airbyte_cdk/sources/declarative/migrations/legacy_to_per_partition_state_migration.py +98 -0
  114. airbyte_cdk/sources/declarative/migrations/state_migration.py +24 -0
  115. airbyte_cdk/sources/declarative/models/__init__.py +2 -0
  116. airbyte_cdk/sources/declarative/models/declarative_component_schema.py +2503 -0
  117. airbyte_cdk/sources/declarative/parsers/__init__.py +3 -0
  118. airbyte_cdk/sources/declarative/parsers/custom_code_compiler.py +157 -0
  119. airbyte_cdk/sources/declarative/parsers/custom_exceptions.py +21 -0
  120. airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py +172 -0
  121. airbyte_cdk/sources/declarative/parsers/manifest_reference_resolver.py +213 -0
  122. airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +3407 -0
  123. airbyte_cdk/sources/declarative/partition_routers/__init__.py +29 -0
  124. airbyte_cdk/sources/declarative/partition_routers/async_job_partition_router.py +65 -0
  125. airbyte_cdk/sources/declarative/partition_routers/cartesian_product_stream_slicer.py +176 -0
  126. airbyte_cdk/sources/declarative/partition_routers/list_partition_router.py +121 -0
  127. airbyte_cdk/sources/declarative/partition_routers/partition_router.py +62 -0
  128. airbyte_cdk/sources/declarative/partition_routers/single_partition_router.py +63 -0
  129. airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +437 -0
  130. airbyte_cdk/sources/declarative/requesters/README.md +56 -0
  131. airbyte_cdk/sources/declarative/requesters/__init__.py +9 -0
  132. airbyte_cdk/sources/declarative/requesters/error_handlers/__init__.py +25 -0
  133. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/__init__.py +23 -0
  134. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/constant_backoff_strategy.py +45 -0
  135. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/exponential_backoff_strategy.py +45 -0
  136. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/header_helper.py +41 -0
  137. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_time_from_header_backoff_strategy.py +70 -0
  138. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_until_time_from_header_backoff_strategy.py +77 -0
  139. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategy.py +17 -0
  140. airbyte_cdk/sources/declarative/requesters/error_handlers/composite_error_handler.py +101 -0
  141. airbyte_cdk/sources/declarative/requesters/error_handlers/default_error_handler.py +147 -0
  142. airbyte_cdk/sources/declarative/requesters/error_handlers/default_http_response_filter.py +40 -0
  143. airbyte_cdk/sources/declarative/requesters/error_handlers/error_handler.py +17 -0
  144. airbyte_cdk/sources/declarative/requesters/error_handlers/http_response_filter.py +179 -0
  145. airbyte_cdk/sources/declarative/requesters/http_job_repository.py +350 -0
  146. airbyte_cdk/sources/declarative/requesters/http_requester.py +433 -0
  147. airbyte_cdk/sources/declarative/requesters/paginators/__init__.py +21 -0
  148. airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py +327 -0
  149. airbyte_cdk/sources/declarative/requesters/paginators/no_pagination.py +76 -0
  150. airbyte_cdk/sources/declarative/requesters/paginators/paginator.py +65 -0
  151. airbyte_cdk/sources/declarative/requesters/paginators/strategies/__init__.py +25 -0
  152. airbyte_cdk/sources/declarative/requesters/paginators/strategies/cursor_pagination_strategy.py +98 -0
  153. airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py +102 -0
  154. airbyte_cdk/sources/declarative/requesters/paginators/strategies/page_increment.py +71 -0
  155. airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py +48 -0
  156. airbyte_cdk/sources/declarative/requesters/paginators/strategies/stop_condition.py +66 -0
  157. airbyte_cdk/sources/declarative/requesters/request_option.py +117 -0
  158. airbyte_cdk/sources/declarative/requesters/request_options/__init__.py +23 -0
  159. airbyte_cdk/sources/declarative/requesters/request_options/datetime_based_request_options_provider.py +92 -0
  160. airbyte_cdk/sources/declarative/requesters/request_options/default_request_options_provider.py +60 -0
  161. airbyte_cdk/sources/declarative/requesters/request_options/interpolated_nested_request_input_provider.py +59 -0
  162. airbyte_cdk/sources/declarative/requesters/request_options/interpolated_request_input_provider.py +68 -0
  163. airbyte_cdk/sources/declarative/requesters/request_options/interpolated_request_options_provider.py +119 -0
  164. airbyte_cdk/sources/declarative/requesters/request_options/request_options_provider.py +79 -0
  165. airbyte_cdk/sources/declarative/requesters/request_path.py +15 -0
  166. airbyte_cdk/sources/declarative/requesters/requester.py +144 -0
  167. airbyte_cdk/sources/declarative/resolvers/__init__.py +41 -0
  168. airbyte_cdk/sources/declarative/resolvers/components_resolver.py +55 -0
  169. airbyte_cdk/sources/declarative/resolvers/config_components_resolver.py +136 -0
  170. airbyte_cdk/sources/declarative/resolvers/http_components_resolver.py +112 -0
  171. airbyte_cdk/sources/declarative/retrievers/__init__.py +19 -0
  172. airbyte_cdk/sources/declarative/retrievers/async_retriever.py +124 -0
  173. airbyte_cdk/sources/declarative/retrievers/file_uploader.py +89 -0
  174. airbyte_cdk/sources/declarative/retrievers/retriever.py +54 -0
  175. airbyte_cdk/sources/declarative/retrievers/simple_retriever.py +702 -0
  176. airbyte_cdk/sources/declarative/schema/__init__.py +25 -0
  177. airbyte_cdk/sources/declarative/schema/default_schema_loader.py +47 -0
  178. airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py +285 -0
  179. airbyte_cdk/sources/declarative/schema/inline_schema_loader.py +19 -0
  180. airbyte_cdk/sources/declarative/schema/json_file_schema_loader.py +92 -0
  181. airbyte_cdk/sources/declarative/schema/schema_loader.py +17 -0
  182. airbyte_cdk/sources/declarative/spec/__init__.py +7 -0
  183. airbyte_cdk/sources/declarative/spec/spec.py +48 -0
  184. airbyte_cdk/sources/declarative/stream_slicers/__init__.py +7 -0
  185. airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py +93 -0
  186. airbyte_cdk/sources/declarative/stream_slicers/stream_slicer.py +25 -0
  187. airbyte_cdk/sources/declarative/transformations/__init__.py +17 -0
  188. airbyte_cdk/sources/declarative/transformations/add_fields.py +146 -0
  189. airbyte_cdk/sources/declarative/transformations/dpath_flatten_fields.py +61 -0
  190. airbyte_cdk/sources/declarative/transformations/flatten_fields.py +52 -0
  191. airbyte_cdk/sources/declarative/transformations/keys_replace_transformation.py +61 -0
  192. airbyte_cdk/sources/declarative/transformations/keys_to_lower_transformation.py +22 -0
  193. airbyte_cdk/sources/declarative/transformations/keys_to_snake_transformation.py +68 -0
  194. airbyte_cdk/sources/declarative/transformations/remove_fields.py +75 -0
  195. airbyte_cdk/sources/declarative/transformations/transformation.py +37 -0
  196. airbyte_cdk/sources/declarative/types.py +25 -0
  197. airbyte_cdk/sources/declarative/yaml_declarative_source.py +67 -0
  198. airbyte_cdk/sources/file_based/README.md +152 -0
  199. airbyte_cdk/sources/file_based/__init__.py +24 -0
  200. airbyte_cdk/sources/file_based/availability_strategy/__init__.py +11 -0
  201. airbyte_cdk/sources/file_based/availability_strategy/abstract_file_based_availability_strategy.py +73 -0
  202. airbyte_cdk/sources/file_based/availability_strategy/default_file_based_availability_strategy.py +149 -0
  203. airbyte_cdk/sources/file_based/config/__init__.py +0 -0
  204. airbyte_cdk/sources/file_based/config/abstract_file_based_spec.py +153 -0
  205. airbyte_cdk/sources/file_based/config/avro_format.py +25 -0
  206. airbyte_cdk/sources/file_based/config/csv_format.py +210 -0
  207. airbyte_cdk/sources/file_based/config/excel_format.py +18 -0
  208. airbyte_cdk/sources/file_based/config/file_based_stream_config.py +99 -0
  209. airbyte_cdk/sources/file_based/config/jsonl_format.py +18 -0
  210. airbyte_cdk/sources/file_based/config/parquet_format.py +25 -0
  211. airbyte_cdk/sources/file_based/config/unstructured_format.py +102 -0
  212. airbyte_cdk/sources/file_based/config/validate_config_transfer_modes.py +81 -0
  213. airbyte_cdk/sources/file_based/discovery_policy/__init__.py +8 -0
  214. airbyte_cdk/sources/file_based/discovery_policy/abstract_discovery_policy.py +21 -0
  215. airbyte_cdk/sources/file_based/discovery_policy/default_discovery_policy.py +33 -0
  216. airbyte_cdk/sources/file_based/exceptions.py +159 -0
  217. airbyte_cdk/sources/file_based/file_based_source.py +466 -0
  218. airbyte_cdk/sources/file_based/file_based_stream_permissions_reader.py +123 -0
  219. airbyte_cdk/sources/file_based/file_based_stream_reader.py +209 -0
  220. airbyte_cdk/sources/file_based/file_record_data.py +22 -0
  221. airbyte_cdk/sources/file_based/file_types/__init__.py +37 -0
  222. airbyte_cdk/sources/file_based/file_types/avro_parser.py +233 -0
  223. airbyte_cdk/sources/file_based/file_types/csv_parser.py +527 -0
  224. airbyte_cdk/sources/file_based/file_types/excel_parser.py +196 -0
  225. airbyte_cdk/sources/file_based/file_types/file_transfer.py +30 -0
  226. airbyte_cdk/sources/file_based/file_types/file_type_parser.py +86 -0
  227. airbyte_cdk/sources/file_based/file_types/jsonl_parser.py +145 -0
  228. airbyte_cdk/sources/file_based/file_types/parquet_parser.py +275 -0
  229. airbyte_cdk/sources/file_based/file_types/unstructured_parser.py +480 -0
  230. airbyte_cdk/sources/file_based/remote_file.py +18 -0
  231. airbyte_cdk/sources/file_based/schema_helpers.py +281 -0
  232. airbyte_cdk/sources/file_based/schema_validation_policies/__init__.py +17 -0
  233. airbyte_cdk/sources/file_based/schema_validation_policies/abstract_schema_validation_policy.py +20 -0
  234. airbyte_cdk/sources/file_based/schema_validation_policies/default_schema_validation_policies.py +52 -0
  235. airbyte_cdk/sources/file_based/stream/__init__.py +13 -0
  236. airbyte_cdk/sources/file_based/stream/abstract_file_based_stream.py +197 -0
  237. airbyte_cdk/sources/file_based/stream/concurrent/__init__.py +0 -0
  238. airbyte_cdk/sources/file_based/stream/concurrent/adapters.py +343 -0
  239. airbyte_cdk/sources/file_based/stream/concurrent/cursor/__init__.py +9 -0
  240. airbyte_cdk/sources/file_based/stream/concurrent/cursor/abstract_concurrent_file_based_cursor.py +59 -0
  241. airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_concurrent_cursor.py +313 -0
  242. airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_final_state_cursor.py +83 -0
  243. airbyte_cdk/sources/file_based/stream/cursor/__init__.py +4 -0
  244. airbyte_cdk/sources/file_based/stream/cursor/abstract_file_based_cursor.py +66 -0
  245. airbyte_cdk/sources/file_based/stream/cursor/default_file_based_cursor.py +149 -0
  246. airbyte_cdk/sources/file_based/stream/default_file_based_stream.py +396 -0
  247. airbyte_cdk/sources/file_based/stream/identities_stream.py +49 -0
  248. airbyte_cdk/sources/file_based/stream/permissions_file_based_stream.py +92 -0
  249. airbyte_cdk/sources/file_based/types.py +10 -0
  250. airbyte_cdk/sources/http_config.py +10 -0
  251. airbyte_cdk/sources/http_logger.py +55 -0
  252. airbyte_cdk/sources/message/__init__.py +19 -0
  253. airbyte_cdk/sources/message/repository.py +137 -0
  254. airbyte_cdk/sources/source.py +95 -0
  255. airbyte_cdk/sources/specs/transfer_modes.py +26 -0
  256. airbyte_cdk/sources/streams/__init__.py +8 -0
  257. airbyte_cdk/sources/streams/availability_strategy.py +84 -0
  258. airbyte_cdk/sources/streams/call_rate.py +704 -0
  259. airbyte_cdk/sources/streams/checkpoint/__init__.py +26 -0
  260. airbyte_cdk/sources/streams/checkpoint/checkpoint_reader.py +335 -0
  261. airbyte_cdk/sources/streams/checkpoint/cursor.py +77 -0
  262. airbyte_cdk/sources/streams/checkpoint/per_partition_key_serializer.py +22 -0
  263. airbyte_cdk/sources/streams/checkpoint/resumable_full_refresh_cursor.py +51 -0
  264. airbyte_cdk/sources/streams/checkpoint/substream_resumable_full_refresh_cursor.py +110 -0
  265. airbyte_cdk/sources/streams/concurrent/README.md +7 -0
  266. airbyte_cdk/sources/streams/concurrent/__init__.py +3 -0
  267. airbyte_cdk/sources/streams/concurrent/abstract_stream.py +96 -0
  268. airbyte_cdk/sources/streams/concurrent/abstract_stream_facade.py +37 -0
  269. airbyte_cdk/sources/streams/concurrent/adapters.py +397 -0
  270. airbyte_cdk/sources/streams/concurrent/availability_strategy.py +94 -0
  271. airbyte_cdk/sources/streams/concurrent/clamping.py +99 -0
  272. airbyte_cdk/sources/streams/concurrent/cursor.py +481 -0
  273. airbyte_cdk/sources/streams/concurrent/cursor_types.py +32 -0
  274. airbyte_cdk/sources/streams/concurrent/default_stream.py +102 -0
  275. airbyte_cdk/sources/streams/concurrent/exceptions.py +18 -0
  276. airbyte_cdk/sources/streams/concurrent/helpers.py +42 -0
  277. airbyte_cdk/sources/streams/concurrent/partition_enqueuer.py +64 -0
  278. airbyte_cdk/sources/streams/concurrent/partition_reader.py +45 -0
  279. airbyte_cdk/sources/streams/concurrent/partitions/__init__.py +3 -0
  280. airbyte_cdk/sources/streams/concurrent/partitions/partition.py +48 -0
  281. airbyte_cdk/sources/streams/concurrent/partitions/partition_generator.py +18 -0
  282. airbyte_cdk/sources/streams/concurrent/partitions/stream_slicer.py +21 -0
  283. airbyte_cdk/sources/streams/concurrent/partitions/types.py +38 -0
  284. airbyte_cdk/sources/streams/concurrent/state_converters/__init__.py +0 -0
  285. airbyte_cdk/sources/streams/concurrent/state_converters/abstract_stream_state_converter.py +182 -0
  286. airbyte_cdk/sources/streams/concurrent/state_converters/datetime_stream_state_converter.py +223 -0
  287. airbyte_cdk/sources/streams/concurrent/state_converters/incrementing_count_stream_state_converter.py +92 -0
  288. airbyte_cdk/sources/streams/core.py +703 -0
  289. airbyte_cdk/sources/streams/http/__init__.py +10 -0
  290. airbyte_cdk/sources/streams/http/availability_strategy.py +54 -0
  291. airbyte_cdk/sources/streams/http/error_handlers/__init__.py +22 -0
  292. airbyte_cdk/sources/streams/http/error_handlers/backoff_strategy.py +28 -0
  293. airbyte_cdk/sources/streams/http/error_handlers/default_backoff_strategy.py +17 -0
  294. airbyte_cdk/sources/streams/http/error_handlers/default_error_mapping.py +86 -0
  295. airbyte_cdk/sources/streams/http/error_handlers/error_handler.py +42 -0
  296. airbyte_cdk/sources/streams/http/error_handlers/error_message_parser.py +19 -0
  297. airbyte_cdk/sources/streams/http/error_handlers/http_status_error_handler.py +110 -0
  298. airbyte_cdk/sources/streams/http/error_handlers/json_error_message_parser.py +52 -0
  299. airbyte_cdk/sources/streams/http/error_handlers/response_models.py +65 -0
  300. airbyte_cdk/sources/streams/http/exceptions.py +61 -0
  301. airbyte_cdk/sources/streams/http/http.py +673 -0
  302. airbyte_cdk/sources/streams/http/http_client.py +531 -0
  303. airbyte_cdk/sources/streams/http/rate_limiting.py +158 -0
  304. airbyte_cdk/sources/streams/http/requests_native_auth/__init__.py +14 -0
  305. airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py +479 -0
  306. airbyte_cdk/sources/streams/http/requests_native_auth/abstract_token.py +34 -0
  307. airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py +436 -0
  308. airbyte_cdk/sources/streams/http/requests_native_auth/token.py +83 -0
  309. airbyte_cdk/sources/streams/permissions/identities_stream.py +75 -0
  310. airbyte_cdk/sources/streams/utils/__init__.py +3 -0
  311. airbyte_cdk/sources/types.py +169 -0
  312. airbyte_cdk/sources/utils/__init__.py +7 -0
  313. airbyte_cdk/sources/utils/casing.py +12 -0
  314. airbyte_cdk/sources/utils/files_directory.py +15 -0
  315. airbyte_cdk/sources/utils/record_helper.py +53 -0
  316. airbyte_cdk/sources/utils/schema_helpers.py +230 -0
  317. airbyte_cdk/sources/utils/slice_logger.py +57 -0
  318. airbyte_cdk/sources/utils/transform.py +277 -0
  319. airbyte_cdk/sources/utils/types.py +7 -0
  320. airbyte_cdk/sql/__init__.py +0 -0
  321. airbyte_cdk/sql/_util/__init__.py +0 -0
  322. airbyte_cdk/sql/_util/hashing.py +34 -0
  323. airbyte_cdk/sql/_util/name_normalizers.py +92 -0
  324. airbyte_cdk/sql/constants.py +32 -0
  325. airbyte_cdk/sql/exceptions.py +235 -0
  326. airbyte_cdk/sql/secrets.py +123 -0
  327. airbyte_cdk/sql/shared/__init__.py +15 -0
  328. airbyte_cdk/sql/shared/catalog_providers.py +145 -0
  329. airbyte_cdk/sql/shared/sql_processor.py +786 -0
  330. airbyte_cdk/sql/types.py +160 -0
  331. airbyte_cdk/test/__init__.py +7 -0
  332. airbyte_cdk/test/catalog_builder.py +81 -0
  333. airbyte_cdk/test/entrypoint_wrapper.py +250 -0
  334. airbyte_cdk/test/mock_http/__init__.py +6 -0
  335. airbyte_cdk/test/mock_http/matcher.py +41 -0
  336. airbyte_cdk/test/mock_http/mocker.py +185 -0
  337. airbyte_cdk/test/mock_http/request.py +103 -0
  338. airbyte_cdk/test/mock_http/response.py +28 -0
  339. airbyte_cdk/test/mock_http/response_builder.py +237 -0
  340. airbyte_cdk/test/state_builder.py +33 -0
  341. airbyte_cdk/test/utils/__init__.py +1 -0
  342. airbyte_cdk/test/utils/data.py +24 -0
  343. airbyte_cdk/test/utils/http_mocking.py +16 -0
  344. airbyte_cdk/test/utils/manifest_only_fixtures.py +59 -0
  345. airbyte_cdk/test/utils/reading.py +26 -0
  346. airbyte_cdk/utils/__init__.py +10 -0
  347. airbyte_cdk/utils/airbyte_secrets_utils.py +80 -0
  348. airbyte_cdk/utils/analytics_message.py +25 -0
  349. airbyte_cdk/utils/constants.py +5 -0
  350. airbyte_cdk/utils/datetime_format_inferrer.py +94 -0
  351. airbyte_cdk/utils/datetime_helpers.py +499 -0
  352. airbyte_cdk/utils/event_timing.py +85 -0
  353. airbyte_cdk/utils/is_cloud_environment.py +18 -0
  354. airbyte_cdk/utils/mapping_helpers.py +162 -0
  355. airbyte_cdk/utils/message_utils.py +26 -0
  356. airbyte_cdk/utils/oneof_option_config.py +33 -0
  357. airbyte_cdk/utils/print_buffer.py +75 -0
  358. airbyte_cdk/utils/schema_inferrer.py +270 -0
  359. airbyte_cdk/utils/slice_hasher.py +37 -0
  360. airbyte_cdk/utils/spec_schema_transformations.py +26 -0
  361. airbyte_cdk/utils/stream_status_utils.py +43 -0
  362. airbyte_cdk/utils/traced_exception.py +145 -0
  363. airbyte_cdk-0.0.0.dev0.dist-info/LICENSE.txt +19 -0
  364. airbyte_cdk-0.0.0.dev0.dist-info/LICENSE_SHORT +1 -0
  365. airbyte_cdk-0.0.0.dev0.dist-info/METADATA +111 -0
  366. airbyte_cdk-0.0.0.dev0.dist-info/RECORD +368 -0
  367. airbyte_cdk-0.0.0.dev0.dist-info/WHEEL +4 -0
  368. 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