airbyte-source-shopify 2.4.14.dev202407181247__py3-none-any.whl → 3.1.0__py3-none-any.whl

source_shopify/shopify_graphql/bulk/job.py

@@ -2,27 +2,38 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import logging
 from dataclasses import dataclass, field
 from datetime import datetime
+from enum import Enum
 from time import sleep, time
 from typing import Any, Final, Iterable, List, Mapping, Optional
 
 import pendulum as pdm
 import requests
-from airbyte_cdk.sources.streams.http import HttpClient
 from requests.exceptions import JSONDecodeError
-from source_shopify.utils import ApiTypeEnum
+from source_shopify.utils import LOGGER, ApiTypeEnum
 from source_shopify.utils import ShopifyRateLimiter as limiter
 
-from ...http_request import ShopifyErrorHandler
+from airbyte_cdk.sources.streams.http import HttpClient
+
 from .exceptions import AirbyteTracedException, ShopifyBulkExceptions
 from .query import ShopifyBulkQuery, ShopifyBulkTemplates
+from .record import ShopifyBulkRecord
 from .retry import bulk_retry_on_exception
 from .status import ShopifyBulkJobStatus
 from .tools import END_OF_FILE, BulkTools
 
 
+class BulkOperationUserErrorCode(Enum):
+    """
+    Possible error codes that can be returned by BulkOperationUserError.
+    https://shopify.dev/docs/api/admin-graphql/latest/enums/BulkOperationUserErrorCode
+    """
+
+    INVALID = "INVALID"
+    OPERATION_IN_PROGRESS = "OPERATION_IN_PROGRESS"
+
+
 @dataclass
 class ShopifyBulkManager:
     http_client: HttpClient
@@ -32,8 +43,8 @@ class ShopifyBulkManager:
     job_size: float
     job_checkpoint_interval: int
 
-    # default logger
-    logger: Final[logging.Logger] = logging.getLogger("airbyte")
+    parent_stream_name: Optional[str] = None
+    parent_stream_cursor: Optional[str] = None
 
     # 10Mb chunk size to save the file
     _retrieve_chunk_size: Final[int] = 1024 * 1024 * 10
@@ -54,7 +65,7 @@ class ShopifyBulkManager:
 
     # currents: _job_id, _job_state, _job_created_at, _job_self_canceled
     _job_id: Optional[str] = field(init=False, default=None)
-    _job_state: str = field(init=False, default=None)  # this string is based on ShopifyBulkJobStatus
+    _job_state: str | None = field(init=False, default=None)  # this string is based on ShopifyBulkJobStatus
     # completed and saved Bulk Job result filename
     _job_result_filename: Optional[str] = field(init=False, default=None)
     # date-time when the Bulk Job was created on the server
@@ -71,8 +82,8 @@ class ShopifyBulkManager:
     _job_last_rec_count: int = field(init=False, default=0)
     # the flag to adjust the next slice from the checkpointed cursor vaue
     _job_adjust_slice_from_checkpoint: bool = field(init=False, default=False)
-    # flag to mark the long running BULK job
-    _job_long_running_cancelation: bool = field(init=False, default=False)
+    # keeps the last checkpointed cursor value for supported streams
+    _job_last_checkpoint_cursor_value: str | None = field(init=False, default=None)
 
     # expand slice factor
     _job_size_expand_factor: int = field(init=False, default=2)
@@ -84,7 +95,7 @@ class ShopifyBulkManager:
     # 2 sec is set as default value to cover the case with the empty-fast-completed jobs
     _job_last_elapsed_time: float = field(init=False, default=2.0)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         self._job_size = self.job_size
         # The upper boundary for slice size is limited by the value from the config, default value is `P30D`
         self._job_size_max = self.job_size
@@ -95,6 +106,8 @@ class ShopifyBulkManager:
         self._job_max_elapsed_time = self.job_termination_threshold
         # how many records should be collected before we use the checkpoining
         self._job_checkpoint_interval = self.job_checkpoint_interval
+        # define Record Producer instance
+        self.record_producer: ShopifyBulkRecord = ShopifyBulkRecord(self.query, self.parent_stream_name, self.parent_stream_cursor)
 
     @property
     def _tools(self) -> BulkTools:
@@ -151,9 +164,20 @@ class ShopifyBulkManager:
         self._job_should_revert_slice = False
         return False
 
+    @property
+    def _supports_checkpointing(self) -> bool:
+        """
+        The flag to determine whether or not the BULK Stream supports the `BULK checkpointing`.
+        """
+        return self.query.supports_checkpointing
+
     @property
     def _job_should_checkpoint(self) -> bool:
-        return self._job_last_rec_count >= self._job_checkpoint_interval
+        return self._supports_checkpointing and self._job_last_rec_count >= self._job_checkpoint_interval
+
+    @property
+    def _job_any_lines_collected(self) -> bool:
+        return self._job_last_rec_count > 0
 
     def _expand_job_size(self) -> None:
         self._job_size += self._job_size_adjusted_expand_factor
@@ -191,18 +215,49 @@ class ShopifyBulkManager:
         # set the running job object count to default
         self._job_last_rec_count = 0
 
+    def _set_checkpointing(self) -> None:
+        # set the flag to adjust the next slice from the checkpointed cursor value
+        self._job_adjust_slice_from_checkpoint = True
+
+    def _reset_checkpointing(self) -> None:
+        # reseting the checkpoint flag, if bulk job has completed normally
+        self._job_adjust_slice_from_checkpoint = False
+
+    def _set_last_checkpoint_cursor_value(self, checkpointed_cursor: str) -> None:
+        """
+        Sets the last checkpoint cursor value.
+
+        Args:
+            checkpointed_cursor (str): The cursor value to set as the last checkpoint. Defaults to None.
+        """
+        self._job_last_checkpoint_cursor_value = checkpointed_cursor
+
+    def _checkpoint_cursor_has_collision(self, checkpointed_cursor: str) -> bool:
+        """
+        Checks if the provided checkpointed cursor collides with the last checkpointed cursor value.
+
+        Args:
+            checkpointed_cursor (str): The cursor value to check for collision. Defaults to None.
+
+        Returns:
+            bool: True if the provided cursor collides with the last checkpointed cursor value, False otherwise.
+        """
+        return self._job_last_checkpoint_cursor_value == checkpointed_cursor
+
     def _job_completed(self) -> bool:
         return self._job_state == ShopifyBulkJobStatus.COMPLETED.value
 
     def _job_canceled(self) -> bool:
         return self._job_state == ShopifyBulkJobStatus.CANCELED.value
 
+    def _job_failed(self) -> bool:
+        return self._job_state == ShopifyBulkJobStatus.FAILED.value
+
     def _job_cancel(self) -> None:
         _, canceled_response = self.http_client.send_request(
             http_method="POST",
             url=self.base_url,
-            data=ShopifyBulkTemplates.cancel(self._job_id),
-            headers={"Content-Type": "application/graphql"},
+            json={"query": ShopifyBulkTemplates.cancel(self._job_id)},
             request_kwargs={},
         )
         # mark the job was self-canceled
@@ -221,21 +276,24 @@ class ShopifyBulkManager:
         else:
             message = f"Elapsed time: {self._job_elapsed_time_in_state} sec"
             if self._job_last_rec_count > 0:
-                count_message = f". Lines collected: {self._job_last_rec_count}"
+                count_message = f". Rows collected: {self._job_last_rec_count}"
                 message = message + count_message
             self._log_state(message)
             self._log_job_msg_count = 0
 
     def _log_state(self, message: Optional[str] = None) -> None:
-        pattern = f"Stream: `{self.http_client._name}`, the BULK Job: `{self._job_id}` is {self._job_state}"
+        pattern = f"Stream: `{self.http_client.name}`, the BULK Job: `{self._job_id}` is {self._job_state}"
         if message:
-            self.logger.info(f"{pattern}. {message}.")
+            LOGGER.info(f"{pattern}. {message}.")
         else:
-            self.logger.info(pattern)
+            LOGGER.info(pattern)
 
     def _job_get_result(self, response: Optional[requests.Response] = None) -> Optional[str]:
         parsed_response = response.json().get("data", {}).get("node", {}) if response else None
-        job_result_url = parsed_response.get("url") if parsed_response else None
+        # get `complete` or `partial` result from collected Bulk Job results
+        full_result_url = parsed_response.get("url") if parsed_response else None
+        partial_result_url = parsed_response.get("partialDataUrl") if parsed_response else None
+        job_result_url = full_result_url if full_result_url else partial_result_url
         if job_result_url:
             # save to local file using chunks to avoid OOM
             filename = self._tools.filename_from_url(job_result_url)
@@ -248,6 +306,13 @@ class ShopifyBulkManager:
                 file.write(END_OF_FILE.encode())
             return filename
 
+    def _job_get_checkpointed_result(self, response: Optional[requests.Response]) -> None:
+        if self._job_any_lines_collected or self._job_should_checkpoint:
+            # set the flag to adjust the next slice from the checkpointed cursor value
+            self._set_checkpointing()
+            # fetch the collected records from CANCELED Job on checkpointing
+            self._job_result_filename = self._job_get_result(response)
+
     def _job_update_state(self, response: Optional[requests.Response] = None) -> None:
         if response:
             self._job_state = response.json().get("data", {}).get("node", {}).get("status")
@@ -257,9 +322,7 @@ class ShopifyBulkManager:
                 self._log_job_state_with_count()
             elif self._job_state in [ShopifyBulkJobStatus.CANCELED.value, ShopifyBulkJobStatus.CANCELING.value]:
                 # do not emit `CANCELED / CANCELING` Bulk Job status, while checkpointing
-                if self._job_should_checkpoint:
-                    pass
-                else:
+                if not self._job_should_checkpoint:
                     self._log_job_state_with_count()
             else:
                 self._log_state()
@@ -273,26 +336,20 @@ class ShopifyBulkManager:
                 f"The BULK Job: `{self._job_id}` exited with {self._job_state}, details: {response.text}"
             )
         else:
-            if self._job_should_checkpoint:
-                # set the flag to adjust the next slice from the checkpointed cursor value
-                self._job_adjust_slice_from_checkpoint = True
-                # fetch the collected records from CANCELED Job on checkpointing
-                self._job_result_filename = self._job_get_result(response)
+            self._job_get_checkpointed_result(response)
 
     def _on_canceling_job(self, **kwargs) -> None:
         sleep(self._job_check_interval)
 
     def _cancel_on_long_running_job(self) -> None:
-        self.logger.info(
-            f"Stream: `{self.http_client._name}` the BULK Job: {self._job_id} runs longer than expected ({self._job_max_elapsed_time} sec). Retry with the reduced `Slice Size` after self-cancelation."
+        LOGGER.info(
+            f"Stream: `{self.http_client.name}` the BULK Job: {self._job_id} runs longer than expected ({self._job_max_elapsed_time} sec). Retry with the reduced `Slice Size` after self-cancelation."
         )
-        self._job_long_running_cancelation = True
         self._job_cancel()
 
     def _cancel_on_checkpointing(self) -> None:
-        self.logger.info(f"Stream: `{self.http_client._name}`, checkpointing after >= `{self._job_checkpoint_interval}` lines collected.")
+        LOGGER.info(f"Stream: `{self.http_client.name}`, checkpointing after >= `{self._job_checkpoint_interval}` rows collected.")
         # set the flag to adjust the next slice from the checkpointed cursor value
-        self._job_adjust_slice_from_checkpoint = True
         self._job_cancel()
 
     def _on_running_job(self, **kwargs) -> None:
@@ -306,10 +363,15 @@ class ShopifyBulkManager:
     def _on_completed_job(self, response: Optional[requests.Response] = None) -> None:
         self._job_result_filename = self._job_get_result(response)
 
-    def _on_failed_job(self, response: requests.Response) -> AirbyteTracedException:
-        raise ShopifyBulkExceptions.BulkJobFailed(
-            f"The BULK Job: `{self._job_id}` exited with {self._job_state}, details: {response.text}",
-        )
+    def _on_failed_job(self, response: requests.Response) -> AirbyteTracedException | None:
+        if not self._supports_checkpointing:
+            raise ShopifyBulkExceptions.BulkJobFailed(
+                f"The BULK Job: `{self._job_id}` exited with {self._job_state}, details: {response.text}",
+            )
+        else:
+            # when the Bulk Job fails, usually there is a `partialDataUrl` available,
+            # we leverage the checkpointing in this case.
+            self._job_get_checkpointed_result(response)
 
     def _on_timeout_job(self, **kwargs) -> AirbyteTracedException:
         raise ShopifyBulkExceptions.BulkJobTimout(
@@ -325,17 +387,17 @@ class ShopifyBulkManager:
         raise ShopifyBulkExceptions.BulkJobError(f"Could not validate the status of the BULK Job `{self._job_id}`. Errors: {errors}.")
 
     def _on_non_handable_job_error(self, errors: List[Mapping[str, Any]]) -> AirbyteTracedException:
-        raise ShopifyBulkExceptions.BulkJobNonHandableError(f"The Stream: `{self.http_client._name}`, Non-handable error occured: {errors}")
+        raise ShopifyBulkExceptions.BulkJobNonHandableError(f"The Stream: `{self.http_client.name}`, Non-handable error occured: {errors}")
 
-    def _get_server_errors(self, response: requests.Response) -> List[Optional[dict]]:
+    def _get_server_errors(self, response: requests.Response) -> List[Optional[Mapping[str, Any]]]:
         server_errors = response.json().get("errors", [])
         return [server_errors] if isinstance(server_errors, str) else server_errors
 
-    def _get_user_errors(self, response: requests.Response) -> List[Optional[dict]]:
+    def _get_user_errors(self, response: requests.Response) -> List[Optional[Mapping[str, Any]]]:
         user_errors = response.json().get("data", {}).get("bulkOperationRunQuery", {}).get("userErrors", [])
         return [user_errors] if isinstance(user_errors, str) else user_errors
 
-    def _collect_bulk_errors(self, response: requests.Response) -> List[Optional[dict]]:
+    def _collect_bulk_errors(self, response: requests.Response) -> List[Optional[Mapping[str, Any]]]:
         try:
             return self._get_server_errors(response) + self._get_user_errors(response)
         except (Exception, JSONDecodeError) as e:
@@ -353,8 +415,7 @@ class ShopifyBulkManager:
         _, response = self.http_client.send_request(
             http_method="POST",
             url=self.base_url,
-            data=ShopifyBulkTemplates.status(self._job_id),
-            headers={"Content-Type": "application/graphql"},
+            json={"query": ShopifyBulkTemplates.status(self._job_id)},
             request_kwargs={},
         )
         self._job_healthcheck(response)
@@ -367,30 +428,23 @@ class ShopifyBulkManager:
         Error example:
         [
             {
+                'code': 'OPERATION_IN_PROGRESS',
                 'field': None,
                 'message': 'A bulk query operation for this app and shop is already in progress: gid://shopify/BulkOperation/4039184154813.',
             }
         ]
         """
-
-        concurrent_job_pattern = "A bulk query operation for this app and shop is already in progress"
         # the errors are handled in `job_job_check_for_errors`
         if errors:
             for error in errors:
-                message = error.get("message", "") if isinstance(error, dict) else ""
-                if concurrent_job_pattern in message:
+                error_code = error.get("code", "") if isinstance(error, dict) else ""
+                if error_code == BulkOperationUserErrorCode.OPERATION_IN_PROGRESS.value:
                     return True
         return False
 
     def _has_reached_max_concurrency(self) -> bool:
         return self._concurrent_attempt == self._concurrent_max_retry
 
-    def _switch_base_url(self) -> None:
-        if self._new_base_url:
-            self.base_url = self._new_base_url
-        else:
-            self.logger.warning(f"Failed switching the `base url`, no `new base url` has been retrieved.")
-
     def _should_switch_shop_name(self, response: requests.Response) -> bool:
         """
         Sometimes the API returns the redirected response that points to the same Store but with different Name:
@@ -402,24 +456,26 @@ class ShopifyBulkManager:
 
         This redirection is related to:
         1) `aliased` or `hidden` store names from being exposed
-        2) migrated to data to the new store, but referenced within the old one stil.
+        2) `migrated` store data to the `new store`, but referenced within the old one stil
 
         reference issue: https://github.com/airbytehq/oncall/issues/5866
         """
         if self.base_url != response.url:
-            self._new_base_url = response.url
+            self.base_url = response.url
             return True
         return False
 
-    @bulk_retry_on_exception(logger)
+    @bulk_retry_on_exception()
     def _job_check_state(self) -> None:
         while not self._job_completed():
             if self._job_canceled():
                 break
+            elif self._job_failed():
+                break
             else:
                 self._job_track_running()
 
-    @bulk_retry_on_exception(logger)
+    @bulk_retry_on_exception()
     def create_job(self, stream_slice: Mapping[str, str], filter_field: str) -> None:
         if stream_slice:
             query = self.query.get(filter_field, stream_slice["start"], stream_slice["end"])
@@ -437,7 +493,7 @@ class ShopifyBulkManager:
         if self._has_running_concurrent_job(errors):
             # when the concurrent job takes place, another job could not be created
             # we typically need to wait and retry, but no longer than 10 min. (see retry in `bulk_retry_on_exception`)
-            raise ShopifyBulkExceptions.BulkJobCreationFailedConcurrentError(f"Failed to create job for stream {self.http_client._name}")
+            raise ShopifyBulkExceptions.BulkJobCreationFailedConcurrentError(f"Failed to create job for stream {self.http_client.name}")
         elif self._should_switch_shop_name(response):
             # assign new shop name, since the one that specified in `config` was redirected to the different one.
             raise ShopifyBulkExceptions.BulkJobRedirectToOtherShopError(f"Switching the `store` name, redirected to: {response.url}")
@@ -459,9 +515,9 @@ class ShopifyBulkManager:
             self._job_id = bulk_response.get("id")
             self._job_created_at = bulk_response.get("createdAt")
             self._job_state = ShopifyBulkJobStatus.CREATED.value
-            self.logger.info(f"Stream: `{self.http_client._name}`, the BULK Job: `{self._job_id}` is {ShopifyBulkJobStatus.CREATED.value}")
+            LOGGER.info(f"Stream: `{self.http_client.name}`, the BULK Job: `{self._job_id}` is {ShopifyBulkJobStatus.CREATED.value}")
 
-    def job_size_normalize(self, start: datetime, end: datetime) -> datetime:
+    def job_size_normalize(self, start: datetime, end: datetime) -> None:
         # adjust slice size when it's bigger than the loop point when it should end,
         # to preserve correct job size adjustments when this is the only job we need to run, based on STATE provided
         requested_slice_size = (end - start).total_days()
@@ -471,9 +527,45 @@ class ShopifyBulkManager:
         step = self._job_size if self._job_size else self._job_size_min
         return slice_start.add(days=step)
 
-    def get_adjusted_job_end(self, slice_start: datetime, slice_end: datetime, checkpointed_cursor: Optional[datetime] = None) -> datetime:
+    def _adjust_slice_end(
+        self, slice_end: datetime, checkpointed_cursor: Optional[str] = None, filter_checkpointed_cursor: Optional[str] = None
+    ) -> datetime:
+        """
+        Choose between the existing `slice_end` value or `checkpointed_cursor` value or `filter_checkpointed_cursor` value, if provided.
+
+        Optionally: raises the `transient` error if the checkpoint collision occurs.
+
+        Note: filter_checkpointed_cursor is only used when cursor field is ID for streams like Customer Address etc.
+        This method should return a datetime from last checkpointed value to adjust slice end, when cursor value is ID (int type)
+        method gets end datetime from filter_checkpointed_cursor, which is value from filter field from last record.
+        See https://github.com/airbytehq/oncall/issues/9052 for more details.
+        """
+
+        if checkpointed_cursor:
+            if self._checkpoint_cursor_has_collision(checkpointed_cursor):
+                raise ShopifyBulkExceptions.BulkJobCheckpointCollisionError(
+                    f"The stream: `{self.http_client.name}` checkpoint collision is detected. Try to increase the `BULK Job checkpoint (rows collected)` to the bigger value. The stream will be synced again during the next sync attempt."
+                )
+            # set the checkpointed cursor value
+            self._set_last_checkpoint_cursor_value(checkpointed_cursor)
+            if isinstance(checkpointed_cursor, str):
+                return pdm.parse(checkpointed_cursor)
+            if isinstance(checkpointed_cursor, int):
+                return pdm.parse(filter_checkpointed_cursor)
+
+        return slice_end
+
+    def get_adjusted_job_end(
+        self,
+        slice_start: datetime,
+        slice_end: datetime,
+        checkpointed_cursor: Optional[str] = None,
+        filter_checkpointed_cursor: Optional[str] = None,
+    ) -> datetime:
         if self._job_adjust_slice_from_checkpoint:
-            return pdm.parse(checkpointed_cursor) if checkpointed_cursor else slice_end
+            # set the checkpointing to default, before the next slice is emitted, to avoid inf.loop
+            self._reset_checkpointing()
+            return self._adjust_slice_end(slice_end, checkpointed_cursor, filter_checkpointed_cursor)
 
         if self._is_long_running_job:
             self._job_size_reduce_next()
@@ -481,8 +573,25 @@ class ShopifyBulkManager:
 
         return slice_end
 
+    def _emit_final_job_message(self, job_current_elapsed_time: int) -> None:
+        final_message = f"Stream: `{self.http_client.name}`, the BULK Job: `{self._job_id}` time elapsed: {job_current_elapsed_time} sec."
+
+        if self._job_any_lines_collected:
+            lines_collected_message = f" Rows collected: {self._job_last_rec_count} --> records: `{self.record_producer.record_composed}`."
+            final_message = final_message + lines_collected_message
+
+        # emit final Bulk job status message
+        LOGGER.info(f"{final_message}")
+
+    def _process_bulk_results(self) -> Iterable[Mapping[str, Any]]:
+        if self._job_result_filename:
+            # produce records from saved bulk job result
+            yield from self.record_producer.read_file(self._job_result_filename)
+        else:
+            yield from []
+
     @limiter.balance_rate_limit(api_type=ApiTypeEnum.graphql.value)
-    def job_check_for_completion(self) -> Optional[str]:
+    def job_get_results(self) -> Optional[Iterable[Mapping[str, Any]]]:
         """
         This method checks the status for the `CREATED` Shopify BULK Job, using it's `ID`.
         The time spent for the Job execution is tracked to understand the effort.
@@ -492,7 +601,7 @@ class ShopifyBulkManager:
         try:
             # track created job until it's COMPLETED
             self._job_check_state()
-            return self._job_result_filename
+            yield from self._process_bulk_results()
         except (
             ShopifyBulkExceptions.BulkJobFailed,
             ShopifyBulkExceptions.BulkJobTimout,
@@ -504,9 +613,8 @@ class ShopifyBulkManager:
             raise bulk_job_error
         finally:
             job_current_elapsed_time = round((time() - job_started), 3)
-            self.logger.info(
-                f"Stream: `{self.http_client._name}`, the BULK Job: `{self._job_id}` time elapsed: {job_current_elapsed_time} sec."
-            )
+            # emit the final Bulk Job log message
+            self._emit_final_job_message(job_current_elapsed_time)
             # check whether or not we should expand or reduce the size of the slice
             self.__adjust_job_size(job_current_elapsed_time)
             # reset the state for COMPLETED job