@@ -372,25 +372,31 @@ def update_auto_restatements(
372372 Args:
373373 next_auto_restatement_ts: A dictionary of snapshot name version to the next auto restatement timestamp.
374374 """
375+ next_auto_restatement_ts_deleted = []
376+ next_auto_restatement_ts_filtered = {}
377+ for k , v in next_auto_restatement_ts .items ():
378+ if v is None :
379+ next_auto_restatement_ts_deleted .append (k )
380+ else :
381+ next_auto_restatement_ts_filtered [k ] = v
382+
375383 for where in snapshot_name_version_filter (
376384 self .engine_adapter ,
377- next_auto_restatement_ts ,
385+ next_auto_restatement_ts_deleted ,
378386 column_prefix = "snapshot" ,
379387 alias = None ,
380388 batch_size = self .SNAPSHOT_BATCH_SIZE ,
381389 ):
382390 self .engine_adapter .delete_from (self .auto_restatements_table , where = where )
383391
384- next_auto_restatement_ts_filtered = {
385- k : v for k , v in next_auto_restatement_ts .items () if v is not None
386- }
387392 if not next_auto_restatement_ts_filtered :
388393 return
389394
390- self .engine_adapter .insert_append (
395+ self .engine_adapter .merge (
391396 self .auto_restatements_table ,
392397 _auto_restatements_to_df (next_auto_restatement_ts_filtered ),
393398 columns_to_types = self ._auto_restatement_columns_to_types ,
399+ unique_key = (exp .column ("snapshot_name" ), exp .column ("snapshot_version" )),
394400 )
395401
396402 def count (self ) -> int :
0 commit comments