1+ import logging
12import os
3+ import time
24from abc import ABC , abstractmethod
35from collections .abc import Iterator , Mapping
46from datetime import date , datetime
57
68import psycopg
79import snowflake .connector
10+ from cryptography .hazmat .backends import default_backend
11+ from cryptography .hazmat .primitives import serialization
812from psycopg .rows import class_row
9- from snowflake .connector import DictCursor
13+ from snowflake .connector import DictCursor , ProgrammingError , SnowflakeConnection
14+ from snowflake .connector .network import ReauthenticationRequest , RetryRequest
1015
11- from constants import DEFAULT_MAX_DATE
1216from model import LoadProgress , T
1317from timer import Timer
1418
1721cursor_fetch_timer = Timer ("cursor_fetch" )
1822transform_timer = Timer ("transform" )
1923
24+ logger = logging .getLogger (__name__ )
25+
2026type DbType = str | float | int | bool | date | datetime
2127
2228
@@ -39,17 +45,21 @@ class Extractor(ABC):
3945 def extract_many (self , cls : type [T ], sql : str , params : dict [str , DbType ]) -> Iterator [list [T ]]:
4046 pass
4147
42- def get_query (self , cls : type [T ], is_historical : bool ) -> str :
43- query = cls .fetch_query (is_historical )
48+ def get_query (self , cls : type [T ], is_historical : bool , start_time : datetime ) -> str :
49+ query = cls .fetch_query (is_historical , start_time )
4450 columns = "," .join (cls .column_aliases ())
4551 columns_raw = "," .join (cls .columns_raw ())
4652 return query .replace ("{COLUMNS}" , columns ).replace ("{COLUMNS_NO_ALIAS}" , columns_raw )
4753
48- def extract_idr_data (self , cls : type [T ], progress : LoadProgress | None ) -> Iterator [list [T ]]:
54+ def extract_idr_data (
55+ self , cls : type [T ], progress : LoadProgress | None , start_time : datetime
56+ ) -> Iterator [list [T ]]:
4957 is_historical = progress is None or progress .is_historical ()
50- fetch_query = self .get_query (cls , is_historical )
58+ fetch_query = self .get_query (cls , is_historical , start_time )
5159 batch_timestamp_col = cls .batch_timestamp_col_alias (is_historical )
5260 update_timestamp_col = cls .update_timestamp_col_alias ()
61+
62+ logger .info ("extracting %s" , cls .table ())
5363 if progress is None :
5464 idr_query_timer .start ()
5565 # No saved progress, process the whole table from the beginning
@@ -64,12 +74,16 @@ def extract_idr_data(self, cls: type[T], progress: LoadProgress | None) -> Itera
6474 idr_query_timer .stop ()
6575 return res
6676
67- previous_batch_complete = progress .batch_completion_ts != DEFAULT_MAX_DATE
68- op = ">" if previous_batch_complete else ">="
77+ 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+
6982 idr_query_timer .start ()
70- # Saved progress found, start processing from where we left
83+ # Saved progress found, start processing from where we left off
7184 update_clause = (
72- f"OR { update_timestamp_col } { update_timestamp_col } { op }
85+ f"""AND ({ update_timestamp_col }
86+ OR { update_timestamp_col }
7387 if update_timestamp_col is not None
7488 else ""
7589 )
@@ -78,13 +92,15 @@ def extract_idr_data(self, cls: type[T], progress: LoadProgress | None) -> Itera
7892 fetch_query .replace (
7993 "{WHERE_CLAUSE}" ,
8094 f"""
81- WHERE
82- ({ update_timestamp_col }
83- AND { batch_timestamp_col } { op } { update_clause }
84- AND { batch_timestamp_col } { op } { get_min_transaction_date ()}
95+ WHERE
96+ (
97+ { batch_timestamp_col }
98+ { update_clause }
99+ )
100+ AND { batch_timestamp_col } { get_min_transaction_date ()}
85101 """ ,
86102 ).replace ("{ORDER_BY}" , f"ORDER BY { batch_timestamp_col } ),
87- {"timestamp" : progress . last_ts },
103+ {"timestamp" : compare_timestamp },
88104 )
89105 idr_query_timer .stop ()
90106 return res
@@ -115,40 +131,67 @@ def extract_single(self, cls: type[T], sql: str, params: dict[str, DbType]) -> T
115131class SnowflakeExtractor (Extractor ):
116132 def __init__ (self , batch_size : int ) -> None :
117133 super ().__init__ ()
118- self .conn = snowflake .connector .connect ( # type: ignore
134+
135+ self .conn = SnowflakeExtractor ._connect ()
136+ self .batch_size = batch_size
137+
138+ @staticmethod
139+ def _connect () -> SnowflakeConnection :
140+ private_key = serialization .load_pem_private_key (
141+ os .environ ["IDR_PRIVATE_KEY" ].encode (), password = None , backend = default_backend ()
142+ )
143+ private_key_bytes = private_key .private_bytes (
144+ encoding = serialization .Encoding .DER ,
145+ format = serialization .PrivateFormat .PKCS8 ,
146+ encryption_algorithm = serialization .NoEncryption (),
147+ )
148+ return snowflake .connector .connect ( # type: ignore
119149 user = os .environ ["IDR_USERNAME" ],
120- password = os . environ [ "IDR_PASSWORD" ] ,
150+ private_key = private_key_bytes ,
121151 account = os .environ ["IDR_ACCOUNT" ],
122152 warehouse = os .environ ["IDR_WAREHOUSE" ],
123153 database = os .environ ["IDR_DATABASE" ],
124154 schema = os .environ ["IDR_SCHEMA" ],
125155 )
126- self .batch_size = batch_size
127156
128157 def extract_many (self , cls : type [T ], sql : str , params : dict [str , DbType ]) -> Iterator [list [T ]]:
129158 cur = None
130- try :
131- cursor_execute_timer .start ()
132- cur = self .conn .cursor (DictCursor )
133- cur .execute (sql , params )
134- cursor_execute_timer .stop ()
135-
136- cursor_fetch_timer .start ()
137- # fetchmany can return list[dict] or list[tuple] but we'll only use
138- # queries that return dicts
139- batch : list [dict [str , DbType ]] = cur .fetchmany (self .batch_size ) # type: ignore[assignment]
140- cursor_fetch_timer .stop ()
141-
142- while len (batch ) > 0 : # type: ignore
143- transform_timer .start ()
144- data = [cls (** {k .lower (): v for k , v in row .items ()}) for row in batch ]
145- transform_timer .stop ()
146-
147- yield data
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 ()
148166
149167 cursor_fetch_timer .start ()
150- batch = cur .fetchmany (self .batch_size ) # type: ignore[assignment]
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]
151171 cursor_fetch_timer .stop ()
152- finally :
153- if cur :
154- cur .close ()
172+
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 ()
0 commit comments