airbyte-cdk 6.5.3rc2__py3-none-any.whl → 6.5.5__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 (198) hide show
  1. airbyte_cdk/__init__.py +17 -2
  2. airbyte_cdk/config_observation.py +10 -3
  3. airbyte_cdk/connector.py +19 -9
  4. airbyte_cdk/connector_builder/connector_builder_handler.py +28 -8
  5. airbyte_cdk/connector_builder/main.py +26 -6
  6. airbyte_cdk/connector_builder/message_grouper.py +95 -25
  7. airbyte_cdk/destinations/destination.py +47 -14
  8. airbyte_cdk/destinations/vector_db_based/config.py +36 -14
  9. airbyte_cdk/destinations/vector_db_based/document_processor.py +49 -11
  10. airbyte_cdk/destinations/vector_db_based/embedder.py +52 -11
  11. airbyte_cdk/destinations/vector_db_based/test_utils.py +14 -4
  12. airbyte_cdk/destinations/vector_db_based/utils.py +8 -2
  13. airbyte_cdk/destinations/vector_db_based/writer.py +15 -4
  14. airbyte_cdk/entrypoint.py +82 -26
  15. airbyte_cdk/exception_handler.py +13 -3
  16. airbyte_cdk/logger.py +10 -2
  17. airbyte_cdk/models/airbyte_protocol.py +11 -5
  18. airbyte_cdk/models/airbyte_protocol_serializers.py +9 -3
  19. airbyte_cdk/models/well_known_types.py +1 -1
  20. airbyte_cdk/sources/abstract_source.py +63 -17
  21. airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py +47 -14
  22. airbyte_cdk/sources/concurrent_source/concurrent_source.py +25 -7
  23. airbyte_cdk/sources/concurrent_source/concurrent_source_adapter.py +27 -6
  24. airbyte_cdk/sources/concurrent_source/thread_pool_manager.py +9 -3
  25. airbyte_cdk/sources/connector_state_manager.py +32 -10
  26. airbyte_cdk/sources/declarative/async_job/job.py +3 -1
  27. airbyte_cdk/sources/declarative/async_job/job_orchestrator.py +68 -14
  28. airbyte_cdk/sources/declarative/async_job/job_tracker.py +24 -6
  29. airbyte_cdk/sources/declarative/async_job/repository.py +3 -1
  30. airbyte_cdk/sources/declarative/auth/declarative_authenticator.py +3 -1
  31. airbyte_cdk/sources/declarative/auth/jwt.py +27 -7
  32. airbyte_cdk/sources/declarative/auth/oauth.py +35 -11
  33. airbyte_cdk/sources/declarative/auth/selective_authenticator.py +3 -1
  34. airbyte_cdk/sources/declarative/auth/token.py +25 -8
  35. airbyte_cdk/sources/declarative/checks/check_stream.py +12 -4
  36. airbyte_cdk/sources/declarative/checks/connection_checker.py +3 -1
  37. airbyte_cdk/sources/declarative/concurrency_level/concurrency_level.py +11 -3
  38. airbyte_cdk/sources/declarative/concurrent_declarative_source.py +106 -50
  39. airbyte_cdk/sources/declarative/datetime/min_max_datetime.py +20 -6
  40. airbyte_cdk/sources/declarative/declarative_source.py +3 -1
  41. airbyte_cdk/sources/declarative/declarative_stream.py +27 -6
  42. airbyte_cdk/sources/declarative/decoders/decoder.py +3 -1
  43. airbyte_cdk/sources/declarative/decoders/json_decoder.py +3 -1
  44. airbyte_cdk/sources/declarative/decoders/pagination_decoder_decorator.py +3 -1
  45. airbyte_cdk/sources/declarative/decoders/xml_decoder.py +6 -2
  46. airbyte_cdk/sources/declarative/extractors/dpath_extractor.py +6 -2
  47. airbyte_cdk/sources/declarative/extractors/record_filter.py +24 -7
  48. airbyte_cdk/sources/declarative/extractors/record_selector.py +10 -3
  49. airbyte_cdk/sources/declarative/extractors/response_to_file_extractor.py +15 -5
  50. airbyte_cdk/sources/declarative/incremental/datetime_based_cursor.py +96 -31
  51. airbyte_cdk/sources/declarative/incremental/global_substream_cursor.py +22 -8
  52. airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py +46 -15
  53. airbyte_cdk/sources/declarative/incremental/per_partition_with_global.py +19 -5
  54. airbyte_cdk/sources/declarative/incremental/resumable_full_refresh_cursor.py +3 -1
  55. airbyte_cdk/sources/declarative/interpolation/interpolated_boolean.py +20 -2
  56. airbyte_cdk/sources/declarative/interpolation/interpolated_mapping.py +5 -1
  57. airbyte_cdk/sources/declarative/interpolation/interpolated_nested_mapping.py +10 -3
  58. airbyte_cdk/sources/declarative/interpolation/interpolated_string.py +6 -2
  59. airbyte_cdk/sources/declarative/interpolation/interpolation.py +7 -1
  60. airbyte_cdk/sources/declarative/interpolation/jinja.py +6 -2
  61. airbyte_cdk/sources/declarative/interpolation/macros.py +19 -4
  62. airbyte_cdk/sources/declarative/manifest_declarative_source.py +106 -24
  63. airbyte_cdk/sources/declarative/migrations/legacy_to_per_partition_state_migration.py +7 -2
  64. airbyte_cdk/sources/declarative/models/declarative_component_schema.py +656 -678
  65. airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py +13 -4
  66. airbyte_cdk/sources/declarative/parsers/manifest_reference_resolver.py +9 -2
  67. airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +782 -232
  68. airbyte_cdk/sources/declarative/partition_routers/cartesian_product_stream_slicer.py +29 -7
  69. airbyte_cdk/sources/declarative/partition_routers/list_partition_router.py +25 -7
  70. airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py +54 -15
  71. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/constant_backoff_strategy.py +6 -2
  72. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/header_helper.py +3 -1
  73. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_time_from_header_backoff_strategy.py +17 -5
  74. airbyte_cdk/sources/declarative/requesters/error_handlers/backoff_strategies/wait_until_time_from_header_backoff_strategy.py +15 -5
  75. airbyte_cdk/sources/declarative/requesters/error_handlers/composite_error_handler.py +3 -1
  76. airbyte_cdk/sources/declarative/requesters/error_handlers/default_error_handler.py +18 -8
  77. airbyte_cdk/sources/declarative/requesters/error_handlers/default_http_response_filter.py +16 -7
  78. airbyte_cdk/sources/declarative/requesters/error_handlers/http_response_filter.py +51 -14
  79. airbyte_cdk/sources/declarative/requesters/http_job_repository.py +29 -8
  80. airbyte_cdk/sources/declarative/requesters/http_requester.py +58 -16
  81. airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py +49 -14
  82. airbyte_cdk/sources/declarative/requesters/paginators/no_pagination.py +3 -1
  83. airbyte_cdk/sources/declarative/requesters/paginators/paginator.py +3 -1
  84. airbyte_cdk/sources/declarative/requesters/paginators/strategies/cursor_pagination_strategy.py +17 -5
  85. airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py +24 -7
  86. airbyte_cdk/sources/declarative/requesters/paginators/strategies/page_increment.py +9 -3
  87. airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py +3 -1
  88. airbyte_cdk/sources/declarative/requesters/paginators/strategies/stop_condition.py +6 -2
  89. airbyte_cdk/sources/declarative/requesters/request_options/datetime_based_request_options_provider.py +19 -6
  90. airbyte_cdk/sources/declarative/requesters/request_options/default_request_options_provider.py +3 -1
  91. airbyte_cdk/sources/declarative/requesters/request_options/interpolated_nested_request_input_provider.py +21 -7
  92. airbyte_cdk/sources/declarative/requesters/request_options/interpolated_request_input_provider.py +18 -6
  93. airbyte_cdk/sources/declarative/requesters/request_options/interpolated_request_options_provider.py +27 -8
  94. airbyte_cdk/sources/declarative/requesters/requester.py +3 -1
  95. airbyte_cdk/sources/declarative/retrievers/async_retriever.py +12 -5
  96. airbyte_cdk/sources/declarative/retrievers/simple_retriever.py +105 -24
  97. airbyte_cdk/sources/declarative/schema/default_schema_loader.py +3 -1
  98. airbyte_cdk/sources/declarative/spec/spec.py +8 -2
  99. airbyte_cdk/sources/declarative/stream_slicers/stream_slicer.py +3 -1
  100. airbyte_cdk/sources/declarative/transformations/add_fields.py +12 -3
  101. airbyte_cdk/sources/declarative/transformations/remove_fields.py +6 -2
  102. airbyte_cdk/sources/declarative/types.py +8 -1
  103. airbyte_cdk/sources/declarative/yaml_declarative_source.py +3 -1
  104. airbyte_cdk/sources/embedded/base_integration.py +14 -4
  105. airbyte_cdk/sources/embedded/catalog.py +16 -4
  106. airbyte_cdk/sources/embedded/runner.py +19 -3
  107. airbyte_cdk/sources/embedded/tools.py +3 -1
  108. airbyte_cdk/sources/file_based/availability_strategy/abstract_file_based_availability_strategy.py +12 -4
  109. airbyte_cdk/sources/file_based/availability_strategy/default_file_based_availability_strategy.py +27 -7
  110. airbyte_cdk/sources/file_based/config/abstract_file_based_spec.py +12 -6
  111. airbyte_cdk/sources/file_based/config/csv_format.py +21 -9
  112. airbyte_cdk/sources/file_based/config/file_based_stream_config.py +6 -2
  113. airbyte_cdk/sources/file_based/config/unstructured_format.py +10 -3
  114. airbyte_cdk/sources/file_based/discovery_policy/abstract_discovery_policy.py +2 -4
  115. airbyte_cdk/sources/file_based/discovery_policy/default_discovery_policy.py +7 -2
  116. airbyte_cdk/sources/file_based/exceptions.py +13 -15
  117. airbyte_cdk/sources/file_based/file_based_source.py +82 -24
  118. airbyte_cdk/sources/file_based/file_based_stream_reader.py +16 -5
  119. airbyte_cdk/sources/file_based/file_types/avro_parser.py +58 -17
  120. airbyte_cdk/sources/file_based/file_types/csv_parser.py +89 -26
  121. airbyte_cdk/sources/file_based/file_types/excel_parser.py +25 -7
  122. airbyte_cdk/sources/file_based/file_types/file_transfer.py +8 -2
  123. airbyte_cdk/sources/file_based/file_types/file_type_parser.py +4 -1
  124. airbyte_cdk/sources/file_based/file_types/jsonl_parser.py +20 -6
  125. airbyte_cdk/sources/file_based/file_types/parquet_parser.py +57 -16
  126. airbyte_cdk/sources/file_based/file_types/unstructured_parser.py +64 -15
  127. airbyte_cdk/sources/file_based/schema_helpers.py +33 -10
  128. airbyte_cdk/sources/file_based/schema_validation_policies/abstract_schema_validation_policy.py +3 -1
  129. airbyte_cdk/sources/file_based/schema_validation_policies/default_schema_validation_policies.py +16 -5
  130. airbyte_cdk/sources/file_based/stream/abstract_file_based_stream.py +33 -10
  131. airbyte_cdk/sources/file_based/stream/concurrent/adapters.py +47 -11
  132. airbyte_cdk/sources/file_based/stream/concurrent/cursor/abstract_concurrent_file_based_cursor.py +13 -22
  133. airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_concurrent_cursor.py +53 -17
  134. airbyte_cdk/sources/file_based/stream/concurrent/cursor/file_based_final_state_cursor.py +17 -5
  135. airbyte_cdk/sources/file_based/stream/cursor/abstract_file_based_cursor.py +3 -1
  136. airbyte_cdk/sources/file_based/stream/cursor/default_file_based_cursor.py +26 -9
  137. airbyte_cdk/sources/file_based/stream/default_file_based_stream.py +67 -21
  138. airbyte_cdk/sources/http_logger.py +5 -1
  139. airbyte_cdk/sources/message/repository.py +18 -4
  140. airbyte_cdk/sources/source.py +17 -7
  141. airbyte_cdk/sources/streams/availability_strategy.py +9 -3
  142. airbyte_cdk/sources/streams/call_rate.py +63 -19
  143. airbyte_cdk/sources/streams/checkpoint/checkpoint_reader.py +31 -7
  144. airbyte_cdk/sources/streams/checkpoint/substream_resumable_full_refresh_cursor.py +6 -2
  145. airbyte_cdk/sources/streams/concurrent/adapters.py +77 -22
  146. airbyte_cdk/sources/streams/concurrent/cursor.py +56 -20
  147. airbyte_cdk/sources/streams/concurrent/default_stream.py +9 -2
  148. airbyte_cdk/sources/streams/concurrent/helpers.py +6 -2
  149. airbyte_cdk/sources/streams/concurrent/partition_enqueuer.py +9 -2
  150. airbyte_cdk/sources/streams/concurrent/partition_reader.py +4 -1
  151. airbyte_cdk/sources/streams/concurrent/partitions/record.py +10 -2
  152. airbyte_cdk/sources/streams/concurrent/partitions/types.py +6 -2
  153. airbyte_cdk/sources/streams/concurrent/state_converters/abstract_stream_state_converter.py +25 -10
  154. airbyte_cdk/sources/streams/concurrent/state_converters/datetime_stream_state_converter.py +32 -16
  155. airbyte_cdk/sources/streams/core.py +77 -22
  156. airbyte_cdk/sources/streams/http/availability_strategy.py +3 -1
  157. airbyte_cdk/sources/streams/http/error_handlers/default_error_mapping.py +4 -1
  158. airbyte_cdk/sources/streams/http/error_handlers/error_handler.py +3 -1
  159. airbyte_cdk/sources/streams/http/error_handlers/http_status_error_handler.py +16 -5
  160. airbyte_cdk/sources/streams/http/error_handlers/response_models.py +9 -3
  161. airbyte_cdk/sources/streams/http/exceptions.py +2 -2
  162. airbyte_cdk/sources/streams/http/http.py +133 -33
  163. airbyte_cdk/sources/streams/http/http_client.py +91 -29
  164. airbyte_cdk/sources/streams/http/rate_limiting.py +23 -7
  165. airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py +19 -6
  166. airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py +38 -11
  167. airbyte_cdk/sources/streams/http/requests_native_auth/token.py +13 -3
  168. airbyte_cdk/sources/types.py +5 -1
  169. airbyte_cdk/sources/utils/record_helper.py +12 -3
  170. airbyte_cdk/sources/utils/schema_helpers.py +9 -3
  171. airbyte_cdk/sources/utils/slice_logger.py +4 -1
  172. airbyte_cdk/sources/utils/transform.py +24 -9
  173. airbyte_cdk/sql/exceptions.py +19 -6
  174. airbyte_cdk/sql/secrets.py +3 -1
  175. airbyte_cdk/sql/shared/catalog_providers.py +13 -4
  176. airbyte_cdk/sql/shared/sql_processor.py +44 -14
  177. airbyte_cdk/test/catalog_builder.py +19 -8
  178. airbyte_cdk/test/entrypoint_wrapper.py +27 -8
  179. airbyte_cdk/test/mock_http/mocker.py +41 -11
  180. airbyte_cdk/test/mock_http/request.py +9 -3
  181. airbyte_cdk/test/mock_http/response.py +3 -1
  182. airbyte_cdk/test/mock_http/response_builder.py +29 -7
  183. airbyte_cdk/test/state_builder.py +10 -2
  184. airbyte_cdk/test/utils/data.py +6 -2
  185. airbyte_cdk/test/utils/http_mocking.py +3 -1
  186. airbyte_cdk/utils/airbyte_secrets_utils.py +3 -1
  187. airbyte_cdk/utils/analytics_message.py +10 -2
  188. airbyte_cdk/utils/datetime_format_inferrer.py +4 -1
  189. airbyte_cdk/utils/mapping_helpers.py +3 -1
  190. airbyte_cdk/utils/message_utils.py +11 -4
  191. airbyte_cdk/utils/print_buffer.py +6 -1
  192. airbyte_cdk/utils/schema_inferrer.py +30 -9
  193. airbyte_cdk/utils/spec_schema_transformations.py +3 -1
  194. airbyte_cdk/utils/traced_exception.py +35 -9
  195. {airbyte_cdk-6.5.3rc2.dist-info → airbyte_cdk-6.5.5.dist-info}/METADATA +7 -6
  196. {airbyte_cdk-6.5.3rc2.dist-info → airbyte_cdk-6.5.5.dist-info}/RECORD +198 -198
  197. {airbyte_cdk-6.5.3rc2.dist-info → airbyte_cdk-6.5.5.dist-info}/LICENSE.txt +0 -0
  198. {airbyte_cdk-6.5.3rc2.dist-info → airbyte_cdk-6.5.5.dist-info}/WHEEL +0 -0
@@ -92,7 +92,9 @@ class AbstractOauth2Authenticator(AuthBase):
92
92
 
93
93
  return payload
94
94
 
95
- def _wrap_refresh_token_exception(self, exception: requests.exceptions.RequestException) -> bool:
95
+ def _wrap_refresh_token_exception(
96
+ self, exception: requests.exceptions.RequestException
97
+ ) -> bool:
96
98
  try:
97
99
  if exception.response is not None:
98
100
  exception_content = exception.response.json()
@@ -102,7 +104,8 @@ class AbstractOauth2Authenticator(AuthBase):
102
104
  return False
103
105
  return (
104
106
  exception.response.status_code in self._refresh_token_error_status_codes
105
- and exception_content.get(self._refresh_token_error_key) in self._refresh_token_error_values
107
+ and exception_content.get(self._refresh_token_error_key)
108
+ in self._refresh_token_error_values
106
109
  )
107
110
 
108
111
  @backoff.on_exception(
@@ -115,14 +118,20 @@ class AbstractOauth2Authenticator(AuthBase):
115
118
  )
116
119
  def _get_refresh_access_token_response(self) -> Any:
117
120
  try:
118
- response = requests.request(method="POST", url=self.get_token_refresh_endpoint(), data=self.build_refresh_request_body())
121
+ response = requests.request(
122
+ method="POST",
123
+ url=self.get_token_refresh_endpoint(),
124
+ data=self.build_refresh_request_body(),
125
+ )
119
126
  if response.ok:
120
127
  response_json = response.json()
121
128
  # Add the access token to the list of secrets so it is replaced before logging the response
122
129
  # An argument could be made to remove the prevous access key from the list of secrets, but unmasking values seems like a security incident waiting to happen...
123
130
  access_key = response_json.get(self.get_access_token_name())
124
131
  if not access_key:
125
- raise Exception("Token refresh API response was missing access token {self.get_access_token_name()}")
132
+ raise Exception(
133
+ "Token refresh API response was missing access token {self.get_access_token_name()}"
134
+ )
126
135
  add_to_secrets(access_key)
127
136
  self._log_response(response)
128
137
  return response_json
@@ -136,7 +145,9 @@ class AbstractOauth2Authenticator(AuthBase):
136
145
  raise DefaultBackoffException(request=e.response.request, response=e.response)
137
146
  if self._wrap_refresh_token_exception(e):
138
147
  message = "Refresh token is invalid or expired. Please re-authenticate from Sources/<your source>/Settings."
139
- raise AirbyteTracedException(internal_message=message, message=message, failure_type=FailureType.config_error)
148
+ raise AirbyteTracedException(
149
+ internal_message=message, message=message, failure_type=FailureType.config_error
150
+ )
140
151
  raise
141
152
  except Exception as e:
142
153
  raise Exception(f"Error while refreshing access token: {e}") from e
@@ -149,7 +160,9 @@ class AbstractOauth2Authenticator(AuthBase):
149
160
  """
150
161
  response_json = self._get_refresh_access_token_response()
151
162
 
152
- return response_json[self.get_access_token_name()], response_json[self.get_expires_in_name()]
163
+ return response_json[self.get_access_token_name()], response_json[
164
+ self.get_expires_in_name()
165
+ ]
153
166
 
154
167
  def _parse_token_expiration_date(self, value: Union[str, int]) -> pendulum.DateTime:
155
168
  """
@@ -6,9 +6,14 @@ from typing import Any, List, Mapping, Optional, Sequence, Tuple, Union
6
6
 
7
7
  import dpath
8
8
  import pendulum
9
- from airbyte_cdk.config_observation import create_connector_config_control_message, emit_configuration_as_airbyte_control_message
9
+ from airbyte_cdk.config_observation import (
10
+ create_connector_config_control_message,
11
+ emit_configuration_as_airbyte_control_message,
12
+ )
10
13
  from airbyte_cdk.sources.message import MessageRepository, NoopMessageRepository
11
- from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_oauth import AbstractOauth2Authenticator
14
+ from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_oauth import (
15
+ AbstractOauth2Authenticator,
16
+ )
12
17
 
13
18
 
14
19
  class Oauth2Authenticator(AbstractOauth2Authenticator):
@@ -50,7 +55,9 @@ class Oauth2Authenticator(AbstractOauth2Authenticator):
50
55
  self._token_expiry_date_format = token_expiry_date_format
51
56
  self._token_expiry_is_time_of_expiration = token_expiry_is_time_of_expiration
52
57
  self._access_token = None
53
- super().__init__(refresh_token_error_status_codes, refresh_token_error_key, refresh_token_error_values)
58
+ super().__init__(
59
+ refresh_token_error_status_codes, refresh_token_error_key, refresh_token_error_values
60
+ )
54
61
 
55
62
  def get_token_refresh_endpoint(self) -> str:
56
63
  return self._token_refresh_endpoint
@@ -153,8 +160,16 @@ class SingleUseRefreshTokenOauth2Authenticator(Oauth2Authenticator):
153
160
  token_expiry_is_time_of_expiration bool: set True it if expires_in is returned as time of expiration instead of the number seconds until expiration
154
161
  message_repository (MessageRepository): the message repository used to emit logs on HTTP requests and control message on config update
155
162
  """
156
- self._client_id = client_id if client_id is not None else dpath.get(connector_config, ("credentials", "client_id"))
157
- self._client_secret = client_secret if client_secret is not None else dpath.get(connector_config, ("credentials", "client_secret"))
163
+ self._client_id = (
164
+ client_id
165
+ if client_id is not None
166
+ else dpath.get(connector_config, ("credentials", "client_id"))
167
+ )
168
+ self._client_secret = (
169
+ client_secret
170
+ if client_secret is not None
171
+ else dpath.get(connector_config, ("credentials", "client_secret"))
172
+ )
158
173
  self._access_token_config_path = access_token_config_path
159
174
  self._refresh_token_config_path = refresh_token_config_path
160
175
  self._token_expiry_date_config_path = token_expiry_date_config_path
@@ -204,18 +219,24 @@ class SingleUseRefreshTokenOauth2Authenticator(Oauth2Authenticator):
204
219
  dpath.new(self._connector_config, self._refresh_token_config_path, new_refresh_token)
205
220
 
206
221
  def get_token_expiry_date(self) -> pendulum.DateTime:
207
- expiry_date = dpath.get(self._connector_config, self._token_expiry_date_config_path, default="")
222
+ expiry_date = dpath.get(
223
+ self._connector_config, self._token_expiry_date_config_path, default=""
224
+ )
208
225
  return pendulum.now().subtract(days=1) if expiry_date == "" else pendulum.parse(expiry_date)
209
226
 
210
227
  def set_token_expiry_date(self, new_token_expiry_date):
211
- dpath.new(self._connector_config, self._token_expiry_date_config_path, str(new_token_expiry_date))
228
+ dpath.new(
229
+ self._connector_config, self._token_expiry_date_config_path, str(new_token_expiry_date)
230
+ )
212
231
 
213
232
  def token_has_expired(self) -> bool:
214
233
  """Returns True if the token is expired"""
215
234
  return pendulum.now("UTC") > self.get_token_expiry_date()
216
235
 
217
236
  @staticmethod
218
- def get_new_token_expiry_date(access_token_expires_in: str, token_expiry_date_format: str = None) -> pendulum.DateTime:
237
+ def get_new_token_expiry_date(
238
+ access_token_expires_in: str, token_expiry_date_format: str = None
239
+ ) -> pendulum.DateTime:
219
240
  if token_expiry_date_format:
220
241
  return pendulum.from_format(access_token_expires_in, token_expiry_date_format)
221
242
  else:
@@ -228,8 +249,12 @@ class SingleUseRefreshTokenOauth2Authenticator(Oauth2Authenticator):
228
249
  str: The current access_token, updated if it was previously expired.
229
250
  """
230
251
  if self.token_has_expired():
231
- new_access_token, access_token_expires_in, new_refresh_token = self.refresh_access_token()
232
- new_token_expiry_date = self.get_new_token_expiry_date(access_token_expires_in, self._token_expiry_date_format)
252
+ new_access_token, access_token_expires_in, new_refresh_token = (
253
+ self.refresh_access_token()
254
+ )
255
+ new_token_expiry_date = self.get_new_token_expiry_date(
256
+ access_token_expires_in, self._token_expiry_date_format
257
+ )
233
258
  self.access_token = new_access_token
234
259
  self.set_refresh_token(new_refresh_token)
235
260
  self.set_token_expiry_date(new_token_expiry_date)
@@ -237,7 +262,9 @@ class SingleUseRefreshTokenOauth2Authenticator(Oauth2Authenticator):
237
262
  # Usually, a class shouldn't care about the implementation details but to keep backward compatibility where we print the
238
263
  # message directly in the console, this is needed
239
264
  if not isinstance(self._message_repository, NoopMessageRepository):
240
- self._message_repository.emit_message(create_connector_config_control_message(self._connector_config))
265
+ self._message_repository.emit_message(
266
+ create_connector_config_control_message(self._connector_config)
267
+ )
241
268
  else:
242
269
  emit_configuration_as_airbyte_control_message(self._connector_config)
243
270
  return self.access_token
@@ -6,7 +6,9 @@ import base64
6
6
  from itertools import cycle
7
7
  from typing import List
8
8
 
9
- from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_token import AbstractHeaderAuthenticator
9
+ from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_token import (
10
+ AbstractHeaderAuthenticator,
11
+ )
10
12
 
11
13
 
12
14
  class MultipleTokenAuthenticator(AbstractHeaderAuthenticator):
@@ -24,7 +26,9 @@ class MultipleTokenAuthenticator(AbstractHeaderAuthenticator):
24
26
  def token(self) -> str:
25
27
  return f"{self._auth_method} {next(self._tokens_iter)}"
26
28
 
27
- def __init__(self, tokens: List[str], auth_method: str = "Bearer", auth_header: str = "Authorization"):
29
+ def __init__(
30
+ self, tokens: List[str], auth_method: str = "Bearer", auth_header: str = "Authorization"
31
+ ):
28
32
  self._auth_method = auth_method
29
33
  self._auth_header = auth_header
30
34
  self._tokens = tokens
@@ -65,7 +69,13 @@ class BasicHttpAuthenticator(AbstractHeaderAuthenticator):
65
69
  def token(self) -> str:
66
70
  return f"{self._auth_method} {self._token}"
67
71
 
68
- def __init__(self, username: str, password: str = "", auth_method: str = "Basic", auth_header: str = "Authorization"):
72
+ def __init__(
73
+ self,
74
+ username: str,
75
+ password: str = "",
76
+ auth_method: str = "Basic",
77
+ auth_header: str = "Authorization",
78
+ ):
69
79
  auth_string = f"{username}:{password}".encode("utf8")
70
80
  b64_encoded = base64.b64encode(auth_string).decode("utf8")
71
81
  self._auth_header = auth_header
@@ -54,7 +54,11 @@ class Record(Mapping[str, Any]):
54
54
 
55
55
  class StreamSlice(Mapping[str, Any]):
56
56
  def __init__(
57
- self, *, partition: Mapping[str, Any], cursor_slice: Mapping[str, Any], extra_fields: Optional[Mapping[str, Any]] = None
57
+ self,
58
+ *,
59
+ partition: Mapping[str, Any],
60
+ cursor_slice: Mapping[str, Any],
61
+ extra_fields: Optional[Mapping[str, Any]] = None,
58
62
  ) -> None:
59
63
  """
60
64
  :param partition: The partition keys representing a unique partition in the stream.
@@ -5,7 +5,12 @@ import time
5
5
  from collections.abc import Mapping as ABCMapping
6
6
  from typing import Any, Mapping, Optional
7
7
 
8
- from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteRecordMessage, AirbyteTraceMessage
8
+ from airbyte_cdk.models import (
9
+ AirbyteLogMessage,
10
+ AirbyteMessage,
11
+ AirbyteRecordMessage,
12
+ AirbyteTraceMessage,
13
+ )
9
14
  from airbyte_cdk.models import Type as MessageType
10
15
  from airbyte_cdk.models.file_transfer_record_message import AirbyteFileTransferRecordMessage
11
16
  from airbyte_cdk.sources.streams.core import StreamData
@@ -32,7 +37,9 @@ def stream_data_to_airbyte_message(
32
37
  # docs/connector-development/cdk-python/schemas.md for details.
33
38
  transformer.transform(data, schema) # type: ignore
34
39
  if is_file_transfer_message:
35
- message = AirbyteFileTransferRecordMessage(stream=stream_name, file=data, emitted_at=now_millis, data={})
40
+ message = AirbyteFileTransferRecordMessage(
41
+ stream=stream_name, file=data, emitted_at=now_millis, data={}
42
+ )
36
43
  else:
37
44
  message = AirbyteRecordMessage(stream=stream_name, data=data, emitted_at=now_millis)
38
45
  return AirbyteMessage(type=MessageType.RECORD, record=message)
@@ -41,4 +48,6 @@ def stream_data_to_airbyte_message(
41
48
  case AirbyteLogMessage():
42
49
  return AirbyteMessage(type=MessageType.LOG, log=data_or_message)
43
50
  case _:
44
- raise ValueError(f"Unexpected type for data_or_message: {type(data_or_message)}: {data_or_message}")
51
+ raise ValueError(
52
+ f"Unexpected type for data_or_message: {type(data_or_message)}: {data_or_message}"
53
+ )
@@ -74,7 +74,9 @@ def _expand_refs(schema: Any, ref_resolver: Optional[RefResolver] = None) -> Non
74
74
  if "$ref" in schema:
75
75
  ref_url = schema.pop("$ref")
76
76
  _, definition = ref_resolver.resolve(ref_url)
77
- _expand_refs(definition, ref_resolver=ref_resolver) # expand refs in definitions as well
77
+ _expand_refs(
78
+ definition, ref_resolver=ref_resolver
79
+ ) # expand refs in definitions as well
78
80
  schema.update(definition)
79
81
  else:
80
82
  for key, value in schema.items():
@@ -152,7 +154,9 @@ class ResourceSchemaLoader:
152
154
  base = os.path.dirname(package.__file__) + "/"
153
155
  else:
154
156
  raise ValueError(f"Package {package} does not have a valid __file__ field")
155
- resolved = jsonref.JsonRef.replace_refs(raw_schema, loader=JsonFileLoader(base, "schemas/shared"), base_uri=base)
157
+ resolved = jsonref.JsonRef.replace_refs(
158
+ raw_schema, loader=JsonFileLoader(base, "schemas/shared"), base_uri=base
159
+ )
156
160
  resolved = resolve_ref_links(resolved)
157
161
  if isinstance(resolved, dict):
158
162
  return resolved
@@ -160,7 +164,9 @@ class ResourceSchemaLoader:
160
164
  raise ValueError(f"Expected resolved to be a dict. Got {resolved}")
161
165
 
162
166
 
163
- def check_config_against_spec_or_exit(config: Mapping[str, Any], spec: ConnectorSpecification) -> None:
167
+ def check_config_against_spec_or_exit(
168
+ config: Mapping[str, Any], spec: ConnectorSpecification
169
+ ) -> None:
164
170
  """
165
171
  Check config object against spec. In case of spec is invalid, throws
166
172
  an exception with validation error description.
@@ -27,7 +27,10 @@ class SliceLogger(ABC):
27
27
  printable_slice = dict(_slice) if _slice else _slice
28
28
  return AirbyteMessage(
29
29
  type=MessageType.LOG,
30
- log=AirbyteLogMessage(level=Level.INFO, message=f"{SliceLogger.SLICE_LOG_PREFIX}{json.dumps(printable_slice, default=str)}"),
30
+ log=AirbyteLogMessage(
31
+ level=Level.INFO,
32
+ message=f"{SliceLogger.SLICE_LOG_PREFIX}{json.dumps(printable_slice, default=str)}",
33
+ ),
31
34
  )
32
35
 
33
36
  @abstractmethod
@@ -9,7 +9,13 @@ from typing import Any, Callable, Dict, Mapping, Optional
9
9
 
10
10
  from jsonschema import Draft7Validator, ValidationError, validators
11
11
 
12
- json_to_python_simple = {"string": str, "number": float, "integer": int, "boolean": bool, "null": type(None)}
12
+ json_to_python_simple = {
13
+ "string": str,
14
+ "number": float,
15
+ "integer": int,
16
+ "boolean": bool,
17
+ "null": type(None),
18
+ }
13
19
  json_to_python = {**json_to_python_simple, **{"object": dict, "array": list}}
14
20
  python_to_json = {v: k for k, v in json_to_python.items()}
15
21
 
@@ -56,9 +62,13 @@ class TypeTransformer:
56
62
  # Do not validate field we do not transform for maximum performance.
57
63
  if key in ["type", "array", "$ref", "properties", "items"]
58
64
  }
59
- self._normalizer = validators.create(meta_schema=Draft7Validator.META_SCHEMA, validators=all_validators)
65
+ self._normalizer = validators.create(
66
+ meta_schema=Draft7Validator.META_SCHEMA, validators=all_validators
67
+ )
60
68
 
61
- def registerCustomTransform(self, normalization_callback: Callable[[Any, Dict[str, Any]], Any]) -> Callable:
69
+ def registerCustomTransform(
70
+ self, normalization_callback: Callable[[Any, Dict[str, Any]], Any]
71
+ ) -> Callable:
62
72
  """
63
73
  Register custom normalization callback.
64
74
  :param normalization_callback function to be used for value
@@ -68,7 +78,9 @@ class TypeTransformer:
68
78
  :return Same callbeck, this is usefull for using registerCustomTransform function as decorator.
69
79
  """
70
80
  if TransformConfig.CustomSchemaNormalization not in self._config:
71
- raise Exception("Please set TransformConfig.CustomSchemaNormalization config before registering custom normalizer")
81
+ raise Exception(
82
+ "Please set TransformConfig.CustomSchemaNormalization config before registering custom normalizer"
83
+ )
72
84
  self._custom_normalizer = normalization_callback
73
85
  return normalization_callback
74
86
 
@@ -120,7 +132,10 @@ class TypeTransformer:
120
132
  return bool(original_item)
121
133
  elif target_type == "array":
122
134
  item_types = set(subschema.get("items", {}).get("type", set()))
123
- if item_types.issubset(json_to_python_simple) and type(original_item) in json_to_python_simple.values():
135
+ if (
136
+ item_types.issubset(json_to_python_simple)
137
+ and type(original_item) in json_to_python_simple.values()
138
+ ):
124
139
  return [original_item]
125
140
  except (ValueError, TypeError):
126
141
  return original_item
@@ -133,7 +148,9 @@ class TypeTransformer:
133
148
  :original_validator: native jsonschema validator callback.
134
149
  """
135
150
 
136
- def normalizator(validator_instance: Callable, property_value: Any, instance: Any, schema: Dict[str, Any]):
151
+ def normalizator(
152
+ validator_instance: Callable, property_value: Any, instance: Any, schema: Dict[str, Any]
153
+ ):
137
154
  """
138
155
  Jsonschema validator callable it uses for validating instance. We
139
156
  override default Draft7Validator to perform value transformation
@@ -191,6 +208,4 @@ class TypeTransformer:
191
208
  def get_error_message(self, e: ValidationError) -> str:
192
209
  instance_json_type = python_to_json[type(e.instance)]
193
210
  key_path = "." + ".".join(map(str, e.path))
194
- return (
195
- f"Failed to transform value {repr(e.instance)} of type '{instance_json_type}' to '{e.validator_value}', key path: '{key_path}'"
196
- )
211
+ return f"Failed to transform value {repr(e.instance)} of type '{instance_json_type}' to '{e.validator_value}', key path: '{key_path}'"
@@ -90,12 +90,18 @@ class AirbyteError(Exception):
90
90
  "original_exception",
91
91
  ]
92
92
  display_properties = {
93
- k: v for k, v in self.__dict__.items() if k not in special_properties and not k.startswith("_") and v is not None
93
+ k: v
94
+ for k, v in self.__dict__.items()
95
+ if k not in special_properties and not k.startswith("_") and v is not None
94
96
  }
95
97
  display_properties.update(self.context or {})
96
- context_str = "\n ".join(f"{str(k).replace('_', ' ').title()}: {v!r}" for k, v in display_properties.items())
98
+ context_str = "\n ".join(
99
+ f"{str(k).replace('_', ' ').title()}: {v!r}" for k, v in display_properties.items()
100
+ )
97
101
  exception_str = (
98
- f"{self.get_message()} ({self.__class__.__name__})" + VERTICAL_SEPARATOR + f"\n{self.__class__.__name__}: {self.get_message()}"
102
+ f"{self.get_message()} ({self.__class__.__name__})"
103
+ + VERTICAL_SEPARATOR
104
+ + f"\n{self.__class__.__name__}: {self.get_message()}"
99
105
  )
100
106
 
101
107
  if self.guidance:
@@ -124,7 +130,9 @@ class AirbyteError(Exception):
124
130
  def __repr__(self) -> str:
125
131
  """Return a string representation of the exception."""
126
132
  class_name = self.__class__.__name__
127
- properties_str = ", ".join(f"{k}={v!r}" for k, v in self.__dict__.items() if not k.startswith("_"))
133
+ properties_str = ", ".join(
134
+ f"{k}={v!r}" for k, v in self.__dict__.items() if not k.startswith("_")
135
+ )
128
136
  return f"{class_name}({properties_str})"
129
137
 
130
138
  def safe_logging_dict(self) -> dict[str, Any]:
@@ -180,7 +188,10 @@ class AirbyteInputError(AirbyteError, ValueError):
180
188
  class AirbyteNameNormalizationError(AirbyteError, ValueError):
181
189
  """Error occurred while normalizing a table or column name."""
182
190
 
183
- guidance = "Please consider renaming the source object if possible, or " "raise an issue in GitHub if not."
191
+ guidance = (
192
+ "Please consider renaming the source object if possible, or "
193
+ "raise an issue in GitHub if not."
194
+ )
184
195
  help_url = NEW_ISSUE_URL
185
196
 
186
197
  raw_name: str | None = None
@@ -205,7 +216,9 @@ class AirbyteConnectorError(AirbyteError):
205
216
  logger = logging.getLogger(f"airbyte.{self.connector_name}")
206
217
 
207
218
  log_paths: list[Path] = [
208
- Path(handler.baseFilename).absolute() for handler in logger.handlers if isinstance(handler, logging.FileHandler)
219
+ Path(handler.baseFilename).absolute()
220
+ for handler in logger.handlers
221
+ if isinstance(handler, logging.FileHandler)
209
222
  ]
210
223
 
211
224
  if log_paths:
@@ -101,7 +101,9 @@ class SecretString(str):
101
101
  handler: GetCoreSchemaHandler,
102
102
  ) -> CoreSchema:
103
103
  """Return a modified core schema for the secret string."""
104
- return core_schema.with_info_after_validator_function(function=cls.validate, schema=handler(str), field_name=handler.field_name)
104
+ return core_schema.with_info_after_validator_function(
105
+ function=cls.validate, schema=handler(str), field_name=handler.field_name
106
+ )
105
107
 
106
108
  @classmethod
107
109
  def __get_pydantic_json_schema__( # noqa: PLW3201 # Pydantic dunder method
@@ -77,13 +77,17 @@ class CatalogProvider:
77
77
  )
78
78
 
79
79
  matching_streams: list[ConfiguredAirbyteStream] = [
80
- stream for stream in self.configured_catalog.streams if stream.stream.name == stream_name
80
+ stream
81
+ for stream in self.configured_catalog.streams
82
+ if stream.stream.name == stream_name
81
83
  ]
82
84
  if not matching_streams:
83
85
  raise exc.AirbyteStreamNotFoundError(
84
86
  stream_name=stream_name,
85
87
  context={
86
- "available_streams": [stream.stream.name for stream in self.configured_catalog.streams],
88
+ "available_streams": [
89
+ stream.stream.name for stream in self.configured_catalog.streams
90
+ ],
87
91
  },
88
92
  )
89
93
 
@@ -121,12 +125,17 @@ class CatalogProvider:
121
125
  if not pks:
122
126
  return []
123
127
 
124
- normalized_pks: list[list[str]] = [[LowerCaseNormalizer.normalize(c) for c in pk] for pk in pks]
128
+ normalized_pks: list[list[str]] = [
129
+ [LowerCaseNormalizer.normalize(c) for c in pk] for pk in pks
130
+ ]
125
131
 
126
132
  for pk_nodes in normalized_pks:
127
133
  if len(pk_nodes) != 1:
128
134
  raise exc.AirbyteError(
129
- message=("Nested primary keys are not supported. " "Each PK column should have exactly one node. "),
135
+ message=(
136
+ "Nested primary keys are not supported. "
137
+ "Each PK column should have exactly one node. "
138
+ ),
130
139
  context={
131
140
  "stream_name": stream_name,
132
141
  "primary_key_nodes": pk_nodes,
@@ -16,7 +16,12 @@ import ulid
16
16
  from airbyte_cdk.sql import exceptions as exc
17
17
  from airbyte_cdk.sql._util.hashing import one_way_hash
18
18
  from airbyte_cdk.sql._util.name_normalizers import LowerCaseNormalizer
19
- from airbyte_cdk.sql.constants import AB_EXTRACTED_AT_COLUMN, AB_META_COLUMN, AB_RAW_ID_COLUMN, DEBUG_MODE
19
+ from airbyte_cdk.sql.constants import (
20
+ AB_EXTRACTED_AT_COLUMN,
21
+ AB_META_COLUMN,
22
+ AB_RAW_ID_COLUMN,
23
+ DEBUG_MODE,
24
+ )
20
25
  from airbyte_cdk.sql.secrets import SecretString
21
26
  from airbyte_cdk.sql.types import SQLTypeConverter
22
27
  from airbyte_protocol_dataclasses.models import AirbyteStateMessage
@@ -100,7 +105,9 @@ class SqlConfig(BaseModel, abc.ABC):
100
105
 
101
106
  Raises `NotImplementedError` if a custom vendor client is not defined.
102
107
  """
103
- raise NotImplementedError(f"The type '{type(self).__name__}' does not define a custom client.")
108
+ raise NotImplementedError(
109
+ f"The type '{type(self).__name__}' does not define a custom client."
110
+ )
104
111
 
105
112
 
106
113
  class SqlProcessorBase(abc.ABC):
@@ -270,7 +277,9 @@ class SqlProcessorBase(abc.ABC):
270
277
  query. To ignore the cache and force a refresh, set 'force_refresh' to True.
271
278
  """
272
279
  if force_refresh and shallow_okay:
273
- raise exc.AirbyteInternalError(message="Cannot force refresh and use shallow query at the same time.")
280
+ raise exc.AirbyteInternalError(
281
+ message="Cannot force refresh and use shallow query at the same time."
282
+ )
274
283
 
275
284
  if force_refresh and table_name in self._cached_table_definitions:
276
285
  self._invalidate_table_cache(table_name)
@@ -315,7 +324,9 @@ class SqlProcessorBase(abc.ABC):
315
324
 
316
325
  if DEBUG_MODE:
317
326
  found_schemas = schemas_list
318
- assert schema_name in found_schemas, f"Schema {schema_name} was not created. Found: {found_schemas}"
327
+ assert (
328
+ schema_name in found_schemas
329
+ ), f"Schema {schema_name} was not created. Found: {found_schemas}"
319
330
 
320
331
  def _quote_identifier(self, identifier: str) -> str:
321
332
  """Return the given identifier, quoted."""
@@ -387,7 +398,8 @@ class SqlProcessorBase(abc.ABC):
387
398
  self._known_schemas_list = [
388
399
  found_schema.split(".")[-1].strip('"')
389
400
  for found_schema in found_schemas
390
- if "." not in found_schema or (found_schema.split(".")[0].lower().strip('"') == database_name.lower())
401
+ if "." not in found_schema
402
+ or (found_schema.split(".")[0].lower().strip('"') == database_name.lower())
391
403
  ]
392
404
  return self._known_schemas_list
393
405
 
@@ -511,7 +523,9 @@ class SqlProcessorBase(abc.ABC):
511
523
  for file_path in files:
512
524
  dataframe = pd.read_json(file_path, lines=True)
513
525
 
514
- sql_column_definitions: dict[str, TypeEngine[Any]] = self._get_sql_column_definitions(stream_name)
526
+ sql_column_definitions: dict[str, TypeEngine[Any]] = self._get_sql_column_definitions(
527
+ stream_name
528
+ )
515
529
 
516
530
  # Remove fields that are not in the schema
517
531
  for col_name in dataframe.columns:
@@ -549,7 +563,10 @@ class SqlProcessorBase(abc.ABC):
549
563
  ) -> None:
550
564
  """Add a column to the given table."""
551
565
  self._execute_sql(
552
- text(f"ALTER TABLE {self._fully_qualified(table.name)} " f"ADD COLUMN {column_name} {column_type}"),
566
+ text(
567
+ f"ALTER TABLE {self._fully_qualified(table.name)} "
568
+ f"ADD COLUMN {column_name} {column_type}"
569
+ ),
553
570
  )
554
571
 
555
572
  def _add_missing_columns_to_table(
@@ -626,8 +643,10 @@ class SqlProcessorBase(abc.ABC):
626
643
  deletion_name = f"{final_table_name}_deleteme"
627
644
  commands = "\n".join(
628
645
  [
629
- f"ALTER TABLE {self._fully_qualified(final_table_name)} RENAME " f"TO {deletion_name};",
630
- f"ALTER TABLE {self._fully_qualified(temp_table_name)} RENAME " f"TO {final_table_name};",
646
+ f"ALTER TABLE {self._fully_qualified(final_table_name)} RENAME "
647
+ f"TO {deletion_name};",
648
+ f"ALTER TABLE {self._fully_qualified(temp_table_name)} RENAME "
649
+ f"TO {final_table_name};",
631
650
  f"DROP TABLE {self._fully_qualified(deletion_name)};",
632
651
  ]
633
652
  )
@@ -646,7 +665,9 @@ class SqlProcessorBase(abc.ABC):
646
665
  """
647
666
  nl = "\n"
648
667
  columns = {self._quote_identifier(c) for c in self._get_sql_column_definitions(stream_name)}
649
- pk_columns = {self._quote_identifier(c) for c in self.catalog_provider.get_primary_keys(stream_name)}
668
+ pk_columns = {
669
+ self._quote_identifier(c) for c in self.catalog_provider.get_primary_keys(stream_name)
670
+ }
650
671
  non_pk_columns = columns - pk_columns
651
672
  join_clause = f"{nl} AND ".join(f"tmp.{pk_col} = final.{pk_col}" for pk_col in pk_columns)
652
673
  set_clause = f"{nl} , ".join(f"{col} = tmp.{col}" for col in non_pk_columns)
@@ -704,16 +725,23 @@ class SqlProcessorBase(abc.ABC):
704
725
  temp_table = self._get_table_by_name(temp_table_name)
705
726
  pk_columns = self.catalog_provider.get_primary_keys(stream_name)
706
727
 
707
- columns_to_update: set[str] = self._get_sql_column_definitions(stream_name=stream_name).keys() - set(pk_columns)
728
+ columns_to_update: set[str] = self._get_sql_column_definitions(
729
+ stream_name=stream_name
730
+ ).keys() - set(pk_columns)
708
731
 
709
732
  # Create a dictionary mapping columns in users_final to users_stage for updating
710
733
  update_values = {
711
- self._get_column_by_name(final_table, column): (self._get_column_by_name(temp_table, column)) for column in columns_to_update
734
+ self._get_column_by_name(final_table, column): (
735
+ self._get_column_by_name(temp_table, column)
736
+ )
737
+ for column in columns_to_update
712
738
  }
713
739
 
714
740
  # Craft the WHERE clause for composite primary keys
715
741
  join_conditions = [
716
- self._get_column_by_name(final_table, pk_column) == self._get_column_by_name(temp_table, pk_column) for pk_column in pk_columns
742
+ self._get_column_by_name(final_table, pk_column)
743
+ == self._get_column_by_name(temp_table, pk_column)
744
+ for pk_column in pk_columns
717
745
  ]
718
746
  join_clause = and_(*join_conditions)
719
747
 
@@ -728,7 +756,9 @@ class SqlProcessorBase(abc.ABC):
728
756
  where_not_exists_clause = self._get_column_by_name(final_table, pk_columns[0]) == null()
729
757
 
730
758
  # Select records from temp_table that are not in final_table
731
- select_new_records_stmt = select(temp_table).select_from(joined_table).where(where_not_exists_clause)
759
+ select_new_records_stmt = (
760
+ select(temp_table).select_from(joined_table).where(where_not_exists_clause)
761
+ )
732
762
 
733
763
  # Craft the INSERT statement using the select statement
734
764
  insert_new_records_stmt = insert(final_table).from_select(