adding state migration for lookbacks

Author: ericvanbenschoten-dv01

airbyte-integrations/connectors/source-stripe/source_stripe/components.py

@@ -1,6 +1,7 @@
 from typing import Iterable, Mapping, Any, Optional, List, Set
-
+from datetime import timedelta, datetime
 from airbyte_cdk.sources.declarative.partition_routers import SubstreamPartitionRouter
+from airbyte_cdk.sources.declarative.migrations.state_migration import StateMigration
 from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources.declarative.types import StreamSlice
 
@@ -21,3 +22,36 @@ def stream_slices(self) -> Iterable[StreamSlice]: # type: ignore
 
         self.logger.info(f"[CustomerIdPartitionRouter] Emitting {len(unique_slices)} unique slices")
         yield from unique_slices
+
+class CustomerBalanceStateMigration(StateMigration):
+    def should_migrate(self, stream_state: Mapping[str, Any]) -> bool:
+        return True
+
+    def migrate(self, stream_state: Mapping[str, Any]) -> Mapping[str, Any]:
+        current_state = stream_state or {}
+        state_created = int(current_state.get("state", {}).get("created", 0))
+        current_parent_state = current_state.get("parent_state", {})
+
+        two_days = int(timedelta(days=2).total_seconds())
+        fourteen_day_floor = int(datetime.utcnow().timestamp()) - int(timedelta(days=14).total_seconds())
+
+        created_customer_balance_transactions = max(state_created - two_days, fourteen_day_floor, 0)
+        updated_customers = max(0, current_parent_state.get("customers", {}).get("updated", 0) - two_days, fourteen_day_floor)
+        updated_invoices = max(0, current_parent_state.get("invoices", {}).get("updated", 0) - two_days, fourteen_day_floor)
+
+        migrated_state = {
+            "state": {
+                "created": created_customer_balance_transactions
+            },
+            "parent_state": {
+                "customers": {
+                    "updated": updated_customers
+                },
+                "invoices": {
+                    "updated": updated_invoices
+                }
+            },
+            "use_global_cursor": True
+        }
+
+        return migrated_state

airbyte-integrations/connectors/source-stripe/source_stripe/manifest.yaml

@@ -308,8 +308,6 @@ definitions:
           $parameters:
             request_parameters:
               types[]: '{{["customer.created", "customer.updated", "customer.deleted"]}}'
-        cursor_field: updated
-        lookback_window_days: 1
 
     subscriptions:
       type: StateDelegatingStream
@@ -354,8 +352,6 @@ definitions:
           $parameters:
             request_parameters:
               types[]: '{{["invoice.created", "invoice.deleted", "invoice.finalization_failed", "invoice.finalized", "invoice.marked_uncollectible", "invoice.overdue", "invoice.paid", "invoice.payment_action_required", "invoice.payment_failed", "invoice.payment_succeeded", "invoice.sent", "invoice.updated", "invoice.voided", "invoice.will_be_due"]}}'
-        cursor_field: updated
-        lookback_window_days: 1
 
     transfers:
       type: StateDelegatingStream
@@ -864,7 +860,7 @@ definitions:
               types[]: '{{["radar.early_fraud_warning.created", "radar.early_fraud_warning.updated"]}}'
 
     customer_balance_transactions:
-      $ref: "#/definitions/entity_stream"
+      $ref: "#/definitions/base_stream"
       $parameters:
         name: customer_balance_transactions
       retriever:
@@ -887,10 +883,12 @@ definitions:
               stream:
                 $ref: "#/definitions/streams/invoices"
               incremental_dependency: true
+      state_migrations:
+        - type: CustomStateMigration
+          class_name: source_stripe.components.CustomerBalanceStateMigration
       incremental_sync:
         $ref: "#/definitions/entity_single_slice_cursor"
         cursor_field: created
-        lookback_window_days: 2
         global_substream_cursor: true
       transformations:
         - type: RemoveFields