#1403 source salesforce: review fixes

Author: davydov-d

airbyte-integrations/connectors/source-salesforce/source_salesforce/streams.py

@@ -187,16 +187,14 @@ def _read_pages(
     ) -> Iterable[StreamData]:
         stream_state = stream_state or {}
         pagination_complete = False
-        next_page_token = None
         records = {}
         next_pages = {}
 
         while not pagination_complete:
             index = 0
             for index, property_chunk in enumerate(self.chunk_properties()):
                 request, response = self._fetch_next_page(stream_slice, stream_state, next_pages.get(index), property_chunk)
-                next_page_token = self.next_page_token(response)
-                next_pages[index] = next_page_token
+                next_pages[index] = self.next_page_token(response)
                 chunk_page_records = records_generator_fn(request, response, stream_state, stream_slice)
                 if not self.too_many_properties:
                     # this is the case when a stream has no primary key
@@ -208,25 +206,28 @@ def _read_pages(
 
                 for record_id, record in chunk_page_records.items():
                     if record_id not in records:
-                        records[record_id] = (record, 0)
+                        records[record_id] = (record, 1)
                         continue
                     incomplete_record, counter = records[record_id]
                     incomplete_record.update(record)
                     counter += 1
                     records[record_id] = (incomplete_record, counter)
 
             for record_id, (record, counter) in records.items():
-                if counter != index:
+                if counter != index + 1:
                     # Because we make multiple calls to query N records (each call to fetch X properties of all the N records),
                     # there's a chance that the number of records corresponding to the query may change between the calls. This
                     # may result in data inconsistency. We skip such records for now and log a warning message.
-                    self.logger.warning(f"Inconsistent record with primary key {record_id} found. Skipping it.")
+                    self.logger.warning(
+                        f"Inconsistent record with primary key {record_id} found. It consists of {counter} chunks instead of {index + 1}. "
+                        f"Skipping it."
+                    )
                     continue
                 yield record
 
             records = {}
 
-            if not next_page_token:
+            if not any(next_pages.values()):
                 pagination_complete = True
 
         # Always return an empty generator just in case no records were ever yielded

airbyte-integrations/connectors/source-salesforce/unit_tests/api_test.py

@@ -625,37 +625,53 @@ def test_rest_stream_init_with_too_many_properties(stream_config, stream_api_v2_
 
 
 def test_too_many_properties(stream_config, stream_api_v2_pk_too_many_properties, requests_mock):
+    stream = generate_stream("Account", stream_config, stream_api_v2_pk_too_many_properties)
+    chunks = len(list(stream.chunk_properties()))
+    assert stream.too_many_properties
+    assert stream.primary_key
+    assert type(stream) == RestSalesforceStream
+    url = next_page_url = "https://fase-account.salesforce.com/services/data/v52.0/queryAll"
     requests_mock.get(
-        "https://fase-account.salesforce.com/services/data/v52.0/queryAll",
+        url,
         [
             {
                 "json": {
-                    "nextRecordsUrl": "next",
+                    "nextRecordsUrl": next_page_url,
                     "records": [{"Id": 1, "propertyA": "A"}, {"Id": 2, "propertyA": "A"}]
                 }
             },
             {
                 "json": {
+                    "nextRecordsUrl": next_page_url,
                     "records": [{"Id": 1, "propertyB": "B"}, {"Id": 2, "propertyB": "B"}]
                 }
             },
+            # 2 for 2 chunks above and 1 for a chunk below
+            *[{"json": {"records": [{"Id": 1}, {"Id": 2}], "nextRecordsUrl": next_page_url}} for _ in range(chunks - 3)],
+            {
+                "json": {
+                    "records": [{"Id": 1}, {"Id": 2}]
+                }
+            },
             {
                 "json": {
-                    "nextRecordsUrl": "next",
                     "records": [{"Id": 3, "propertyA": "A"}, {"Id": 4, "propertyA": "A"}]
                 }
             },
             {
                 "json": {
                     "records": [{"Id": 3, "propertyB": "B"}, {"Id": 4, "propertyB": "B"}]
                 }
+            },
+            # 2 for 2 chunks above and 1 for a chunk below
+            *[{"json": {"records": [{"Id": 3}, {"Id": 4}]}} for _ in range(chunks - 3)],
+            {
+                "json": {
+                    "records": [{"Id": 3}, {"Id": 4}]
+                }
             }
         ]
     )
-    stream = generate_stream("Account", stream_config, stream_api_v2_pk_too_many_properties)
-    assert stream.too_many_properties
-    assert stream.primary_key
-    assert type(stream) == RestSalesforceStream
     records = list(stream.read_records(sync_mode=SyncMode.full_refresh))
     assert records == [
         {"Id": 1, "propertyA": "A", "propertyB": "B"},