From f8f547670ab41f54af21071b7c2dddd32afecf14 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9Crdeshmukh15=E2=80=9D?= <rdeshmukh@talend.com>
Date: Fri, 17 May 2024 14:36:32 +0000
Subject: [PATCH 1/2] changes to payments stream implementation to fetch the
 data for all location_id

---
 tap_square/client.py  | 39 +++++++++++++++++++++++++--------------
 tap_square/streams.py |  4 +++-
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/tap_square/client.py b/tap_square/client.py
index bfcb5df..53872f2 100644
--- a/tap_square/client.py
+++ b/tap_square/client.py
@@ -253,23 +253,34 @@ def get_refunds(self, start_time, bookmarked_cursor):
             body,
             'refunds')
 
-    def get_payments(self, start_time, bookmarked_cursor):
-        start_time = utils.strptime_to_utc(start_time)
-        start_time = start_time - timedelta(milliseconds=1)
-        start_time = utils.strftime(start_time)
+    def get_payments(self, location_id, start_time, bookmarked_cursor):
+        if bookmarked_cursor:
+            cursor = bookmarked_cursor
+        else:
+            cursor = '__initial__' # initial value so while loop is always entered one time
 
-        body = {
-        }
-        body['begin_time'] = start_time
+        end_time = utils.strftime(utils.now(), utils.DATETIME_PARSE)
+        while cursor:
+            if cursor == '__initial__':
+                # initial text was needed to go into the while loop, but api needs
+                # it to be a valid bookmarked cursor or None
+                cursor = bookmarked_cursor
 
-        if bookmarked_cursor:
-            body['cursor'] = bookmarked_cursor
+            with singer.http_request_timer('GET payments'):
+                result = self._retryable_v2_method(
+                    lambda bdy: self._client.payments.list_payments(
+                        location_id=location_id,
+                        begin_time=start_time,
+                        end_time=end_time,
+                        cursor=cursor,
+                        limit=1000,
+                    ),
+                    None,
+                )
 
-        yield from self._get_v2_objects(
-            'payments',
-            lambda bdy: self._client.payments.list_payments(**bdy),
-            body,
-            'payments')
+            yield (result.body.get('payments', []), result.body.get('cursor'))
+
+            cursor = result.body.get('cursor')
 
     def get_cash_drawer_shifts(self, location_id, start_time, bookmarked_cursor):
         if bookmarked_cursor:
diff --git a/tap_square/streams.py b/tap_square/streams.py
index 260b2e6..358594c 100644
--- a/tap_square/streams.py
+++ b/tap_square/streams.py
@@ -189,7 +189,9 @@ class Payments(FullTableStream):
     object_type = 'PAYMENT'
 
     def get_pages(self, bookmarked_cursor, start_time):
-        yield from self.client.get_payments(start_time, bookmarked_cursor)
+        for location_id in Locations.get_all_location_ids(self.client):
+            # Payments requests can only take up to 1 location_id at a time
+            yield from self.client.get_payments(location_id, start_time, bookmarked_cursor)
 
 
 class Orders(Stream):

From b95bda0379a5984a70483d2becef1fe9dc2d4b5a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9Crdeshmukh15=E2=80=9D?= <rdeshmukh@talend.com>
Date: Fri, 24 May 2024 13:51:39 +0000
Subject: [PATCH 2/2] removes unnecessary write_state statements

---
 tap_square/streams.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tap_square/streams.py b/tap_square/streams.py
index 358594c..f9fd77a 100644
--- a/tap_square/streams.py
+++ b/tap_square/streams.py
@@ -82,8 +82,6 @@ def sync(self, state, stream_schema, stream_metadata, config, transformer):
                     self.tap_stream_id,
                     transformed_record,
                 )
-            singer.write_bookmark(state, self.tap_stream_id, 'cursor', cursor)
-            singer.write_state(state)
 
         state = singer.clear_bookmark(state, self.tap_stream_id, 'cursor')
         singer.write_state(state)