11import logging
22import os
3- import time
43from abc import ABC , abstractmethod
54from collections .abc import Iterator , Mapping
6- from datetime import date , datetime
5+ from datetime import UTC , date , datetime
76
87import psycopg
98import snowflake .connector
109from cryptography .hazmat .backends import default_backend
1110from cryptography .hazmat .primitives import serialization
1211from psycopg .rows import class_row
13- from snowflake .connector import DictCursor , ProgrammingError , SnowflakeConnection
14- from snowflake .connector .network import ReauthenticationRequest , RetryRequest
12+ from snowflake .connector import DictCursor , SnowflakeConnection
1513
14+ from constants import DEFAULT_MIN_DATE
1615from model import LoadProgress , T
1716from timer import Timer
1817
@@ -33,18 +32,25 @@ def print_timers() -> None:
3332 transform_timer .print_results ()
3433
3534
36- def get_min_transaction_date () -> str :
35+ def get_min_transaction_date () -> datetime :
3736 min_date = os .environ .get ("PIPELINE_MIN_TRANSACTION_DATE" )
3837 if min_date is not None :
39- return min_date
40- return "0001-01-01"
38+ return datetime . strptime ( min_date , "%Y-%m-%d" ). replace ( tzinfo = UTC )
39+ return datetime . strptime ( "0001-01-01" , "%Y-%m-%d" ). replace ( tzinfo = UTC )
4140
4241
4342class Extractor (ABC ):
4443 @abstractmethod
4544 def extract_many (self , cls : type [T ], sql : str , params : dict [str , DbType ]) -> Iterator [list [T ]]:
4645 pass
4746
47+ @abstractmethod
48+ def reconnect (self ) -> None :
49+ pass
50+
51+ def _greatest_col (self , cols : list [str ]) -> str :
52+ return f"GREATEST({ ',' .join (cols )}
53+
4854 def get_query (self , cls : type [T ], is_historical : bool , start_time : datetime ) -> str :
4955 query = cls .fetch_query (is_historical , start_time )
5056 columns = "," .join (cls .column_aliases ())
@@ -56,9 +62,15 @@ def extract_idr_data(
5662 ) -> Iterator [list [T ]]:
5763 is_historical = progress is None or progress .is_historical ()
5864 fetch_query = self .get_query (cls , is_historical , start_time )
59- batch_timestamp_col = cls .batch_timestamp_col_alias (is_historical )
60- update_timestamp_col = cls .update_timestamp_col_alias ()
61-
65+ batch_timestamp_cols = cls .batch_timestamp_col_alias (is_historical )
66+ # GREATEST doesn't work with nulls so we need to coalesce here
67+ update_timestamp_cols = [
68+ f"COALESCE({ col } { DEFAULT_MIN_DATE } for col in cls .update_timestamp_col_alias ()
69+ ]
70+ # We need to create batches using the most recent timestamp from all of the
71+ # insert/update timestamps
72+ batch_timestamp_clause = self ._greatest_col ([* batch_timestamp_cols , * update_timestamp_cols ])
73+ min_transaction_date = get_min_transaction_date ()
6274 logger .info ("extracting %s" , cls .table ())
6375 if progress is None :
6476 idr_query_timer .start ()
@@ -67,39 +79,31 @@ def extract_idr_data(
6779 cls ,
6880 fetch_query .replace (
6981 "{WHERE_CLAUSE}" ,
70- f"WHERE { batch_timestamp_col } { get_min_transaction_date () } ,
71- ).replace ("{ORDER_BY}" , f"ORDER BY { batch_timestamp_col } ),
82+ f"WHERE { batch_timestamp_clause } { min_transaction_date } ,
83+ ).replace ("{ORDER_BY}" , f"ORDER BY { batch_timestamp_clause } ),
7284 {},
7385 )
7486 idr_query_timer .stop ()
7587 return res
7688
7789 previous_batch_complete = progress .batch_complete_ts >= progress .batch_start_ts
78- logger .info ("previous batch complete: %s" , previous_batch_complete )
79-
80- compare_timestamp = progress .batch_start_ts if previous_batch_complete else progress .last_ts
81-
90+ # If we've completed the last batch, there shouldn't be any additional records
91+ # with the same timestamp
92+ op = ">" if previous_batch_complete else ">="
93+ # insertion timestamps aren't always representative of the time the data is available in
94+ # Snowflake, so we should always start loading from the most recent timestamp
95+ # that we've already fetched
96+ compare_timestamp = max (min_transaction_date , progress .last_ts )
8297 idr_query_timer .start ()
8398 # Saved progress found, start processing from where we left off
84- update_clause = (
85- f"""AND ({ update_timestamp_col }
86- OR { update_timestamp_col }
87- if update_timestamp_col is not None
88- else ""
89- )
9099 res = self .extract_many (
91100 cls ,
92101 fetch_query .replace (
93102 "{WHERE_CLAUSE}" ,
94103 f"""
95- WHERE
96- (
97- { batch_timestamp_col }
98- { update_clause }
99- )
100- AND { batch_timestamp_col } { get_min_transaction_date ()}
104+ WHERE { batch_timestamp_clause } { op }
101105 """ ,
102- ).replace ("{ORDER_BY}" , f"ORDER BY { batch_timestamp_col } ),
106+ ).replace ("{ORDER_BY}" , f"ORDER BY { batch_timestamp_clause } ),
103107 {"timestamp" : compare_timestamp },
104108 )
105109 idr_query_timer .stop ()
@@ -109,9 +113,13 @@ def extract_idr_data(
109113class PostgresExtractor (Extractor ):
110114 def __init__ (self , connection_string : str , batch_size : int ) -> None :
111115 super ().__init__ ()
116+ self .connection_string = connection_string
112117 self .conn = psycopg .connect (connection_string )
113118 self .batch_size = batch_size
114119
120+ def reconnect (self ) -> None :
121+ self .conn = psycopg .connect (self .connection_string )
122+
115123 def extract_many (
116124 self , cls : type [T ], sql : str , params : Mapping [str , DbType ]
117125 ) -> Iterator [list [T ]]:
@@ -135,6 +143,9 @@ def __init__(self, batch_size: int) -> None:
135143 self .conn = SnowflakeExtractor ._connect ()
136144 self .batch_size = batch_size
137145
146+ def reconnect (self ) -> None :
147+ SnowflakeExtractor ._connect ()
148+
138149 @staticmethod
139150 def _connect () -> SnowflakeConnection :
140151 private_key = serialization .load_pem_private_key (
@@ -156,42 +167,31 @@ def _connect() -> SnowflakeConnection:
156167
157168 def extract_many (self , cls : type [T ], sql : str , params : dict [str , DbType ]) -> Iterator [list [T ]]:
158169 cur = None
159- max_attempts = 5
160- for attempt in range (max_attempts ):
161- try :
162- cursor_execute_timer .start ()
163- cur = self .conn .cursor (DictCursor )
164- cur .execute (sql , params )
165- cursor_execute_timer .stop ()
170+
171+ try :
172+ cursor_execute_timer .start ()
173+ cur = self .conn .cursor (DictCursor )
174+ cur .execute (sql , params )
175+ cursor_execute_timer .stop ()
176+
177+ cursor_fetch_timer .start ()
178+ # fetchmany can return list[dict] or list[tuple] but we'll only use
179+ # queries that return dicts
180+ batch : list [dict [str , DbType ]] = cur .fetchmany (self .batch_size ) # type: ignore[assignment]
181+ cursor_fetch_timer .stop ()
182+
183+ while len (batch ) > 0 : # type: ignore
184+ transform_timer .start ()
185+ data = [cls (** {k .lower (): v for k , v in row .items ()}) for row in batch ]
186+ transform_timer .stop ()
187+
188+ yield data
166189
167190 cursor_fetch_timer .start ()
168- # fetchmany can return list[dict] or list[tuple] but we'll only use
169- # queries that return dicts
170- batch : list [dict [str , DbType ]] = cur .fetchmany (self .batch_size ) # type: ignore[assignment]
191+ batch = cur .fetchmany (self .batch_size ) # type: ignore[assignment]
171192 cursor_fetch_timer .stop ()
193+ return
172194
173- while len (batch ) > 0 : # type: ignore
174- transform_timer .start ()
175- data = [cls (** {k .lower (): v for k , v in row .items ()}) for row in batch ]
176- transform_timer .stop ()
177-
178- yield data
179-
180- cursor_fetch_timer .start ()
181- batch = cur .fetchmany (self .batch_size ) # type: ignore[assignment]
182- cursor_fetch_timer .stop ()
183- return
184- # Snowflake will throw a reauth error if the pipeline has been running for several hours
185- # but it seems to be wrapped in a ProgrammingError.
186- # Unclear the best way to handle this, it will require a bit more trial and error
187- except (ReauthenticationRequest , RetryRequest , ProgrammingError ) as ex :
188- logger .warning ("received transient error, retrying..." , exc_info = ex )
189- if attempt == max_attempts - 1 :
190- logger .error ("max attempts exceeded" )
191- raise ex
192- self .conn = SnowflakeExtractor ._connect ()
193- time .sleep (1 )
194-
195- finally :
196- if cur :
197- cur .close ()
195+ finally :
196+ if cur :
197+ cur .close ()
0 commit comments