1
1
from typing import Iterable , Mapping , Any , Optional , List , Set
2
-
2
+ from datetime import timedelta , datetime
3
3
from airbyte_cdk .sources .declarative .partition_routers import SubstreamPartitionRouter
4
+ from airbyte_cdk .sources .declarative .migrations .state_migration import StateMigration
4
5
from airbyte_cdk .models import SyncMode
5
6
from airbyte_cdk .sources .declarative .types import StreamSlice
6
7
@@ -21,3 +22,36 @@ def stream_slices(self) -> Iterable[StreamSlice]: # type: ignore
21
22
22
23
self .logger .info (f"[CustomerIdPartitionRouter] Emitting { len (unique_slices )} unique slices" )
23
24
yield from unique_slices
25
+
26
+ class CustomerBalanceStateMigration (StateMigration ):
27
+ def should_migrate (self , stream_state : Mapping [str , Any ]) -> bool :
28
+ return True
29
+
30
+ def migrate (self , stream_state : Mapping [str , Any ]) -> Mapping [str , Any ]:
31
+ current_state = stream_state or {}
32
+ state_created = int (current_state .get ("state" , {}).get ("created" , 0 ))
33
+ current_parent_state = current_state .get ("parent_state" , {})
34
+
35
+ two_days = int (timedelta (days = 2 ).total_seconds ())
36
+ fourteen_day_floor = int (datetime .utcnow ().timestamp ()) - int (timedelta (days = 14 ).total_seconds ())
37
+
38
+ created_customer_balance_transactions = max (state_created - two_days , fourteen_day_floor , 0 )
39
+ updated_customers = max (0 , current_parent_state .get ("customers" , {}).get ("updated" , 0 ) - two_days , fourteen_day_floor )
40
+ updated_invoices = max (0 , current_parent_state .get ("invoices" , {}).get ("updated" , 0 ) - two_days , fourteen_day_floor )
41
+
42
+ migrated_state = {
43
+ "state" : {
44
+ "created" : created_customer_balance_transactions
45
+ },
46
+ "parent_state" : {
47
+ "customers" : {
48
+ "updated" : updated_customers
49
+ },
50
+ "invoices" : {
51
+ "updated" : updated_invoices
52
+ }
53
+ },
54
+ "use_global_cursor" : True
55
+ }
56
+
57
+ return migrated_state
0 commit comments