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,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()