Skip to content

Commit 359db5c

Browse files
committed
Output record count metric from batch files insert
1 parent 2b69ad8 commit 359db5c

File tree

2 files changed

+9
-4
lines changed

2 files changed

+9
-4
lines changed

Diff for: target_snowflake/connector.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -553,7 +553,8 @@ def merge_from_stage(
553553
key_properties=key_properties,
554554
)
555555
self.logger.debug("Merging with SQL: %s", merge_statement)
556-
conn.execute(merge_statement, **kwargs)
556+
result = conn.execute(merge_statement, **kwargs)
557+
return result.rowcount
557558

558559
def copy_from_stage(
559560
self,
@@ -578,7 +579,8 @@ def copy_from_stage(
578579
file_format=file_format,
579580
)
580581
self.logger.debug("Copying with SQL: %s", copy_statement)
581-
conn.execute(copy_statement, **kwargs)
582+
result = conn.execute(copy_statement, **kwargs)
583+
return result.rowcount
582584

583585
def drop_file_format(self, file_format: str) -> None:
584586
"""Drop a file format in the schema.

Diff for: target_snowflake/sinks.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ def insert_batch_files_via_internal_stage(
190190

191191
if self.key_properties:
192192
# merge into destination table
193-
self.connector.merge_from_stage(
193+
record_count = self.connector.merge_from_stage(
194194
full_table_name=full_table_name,
195195
schema=self.schema,
196196
sync_id=sync_id,
@@ -199,13 +199,16 @@ def insert_batch_files_via_internal_stage(
199199
)
200200

201201
else:
202-
self.connector.copy_from_stage(
202+
record_count = self.connector.copy_from_stage(
203203
full_table_name=full_table_name,
204204
schema=self.schema,
205205
sync_id=sync_id,
206206
file_format=file_format,
207207
)
208208

209+
with self.record_counter_metric as counter:
210+
counter.increment(record_count)
211+
209212
finally:
210213
self.logger.debug("Cleaning up after batch processing")
211214
self.connector.drop_file_format(file_format=file_format)

0 commit comments

Comments
 (0)