from __future__ import annotations import re from collections.abc import Sequence from dataclasses import dataclass, field from itertools import chain from typing import Any from config import OPERATION_MODE_STANDARD, S3_CSV_READ_MODE_BUFFERED, ConnectorConfig from csv_parser import iter_csv_rows from csv_parser import new as new_csv_parser from errors import SnapshotError from fivetran.operations import OperationSink, TimedOperationSink, raw_csv_record_writer_factory from fivetran.sdk_records import DiagnosticColumnSource, build_raw_csv_operation_writers from metadata import TableName from observability import OperationStats, elapsed_ms, info, info_with_memory, timer_start, warning, warning_with_memory from pipeline import DecodedSink, RawSink, SampledSink from row_decoder import CSVRowDecoder from snapshot.schema import ( SNAPSHOT_PREFIX, SnapshotTable, snapshot_metadata_pos, ) from snapshot.tables import iter_snapshot_tables from state import ( SNAPSHOT_CURRENT_FILE, SNAPSHOT_NEXT_ROW, CheckpointPolicy, ConnectorState, finish_table_resync_if_complete, mark_table_resync_truncated, set_snapshot_tso, snapshots_complete, table_resync_names, table_resync_truncate_done, ) from state import ( snapshot_tso as state_snapshot_tso, ) from storage.csv_file import open_csv_text from storage.object_store import ObjectStore, join_key from type_mapping import is_source_type_syncable _SNAPSHOT_STAGE_SAMPLE_INTERVAL_ROWS = 4096 def _snapshot_files(store: ObjectStore, table: SnapshotTable, snapshot_prefix: str = SNAPSHOT_PREFIX) -> list[str]: marker = join_key(snapshot_prefix, f"{table.snapshot_file_stem}.") return sorted( (key for key in store.iter_keys(marker) if key.startswith(marker) and key.endswith(".csv")), key=_natural_sort_key, ) def _natural_sort_key(value: str) -> list[int | str]: return [int(part) if part.isdigit() else part for part in re.split(r"(\d+)", value)] def _unsupported_primary_key_columns(table: SnapshotTable) -> list[str]: primary_key = set(table.primary_key) return [ column.name for column in table.columns if column.name in primary_key and not is_source_type_syncable(column.source_type, precision=column.precision) ] @dataclass class SnapshotStageStats: read_duration_ms: float = 0.0 decode_duration_ms: float = 0.0 write_duration_ms: float = 0.0 read_row_count: int = 0 decode_row_count: int = 0 write_row_count: int = 0 decode_sample_count: int = 0 write_sample_count: int = 0 def merge(self, other: SnapshotStageStats) -> None: self.read_duration_ms += other.read_duration_ms self.decode_duration_ms += other.decode_duration_ms self.write_duration_ms += other.write_duration_ms self.read_row_count += other.read_row_count self.decode_row_count += other.decode_row_count self.write_row_count += other.write_row_count self.decode_sample_count += other.decode_sample_count self.write_sample_count += other.write_sample_count def as_log_fields(self) -> dict[str, float | int]: read_parse_duration_ms = round(self.read_duration_ms + self.decode_duration_ms, 3) return { "read_parse_duration_ms": read_parse_duration_ms, "read_duration_ms": round(self.read_duration_ms, 3), "read_row_count": self.read_row_count, "read_us_per_row": _us_per_row(self.read_duration_ms, self.read_row_count), "decode_duration_ms": round(self.decode_duration_ms, 3), "decode_row_count": self.decode_row_count, "decode_sample_count": self.decode_sample_count, "decode_us_per_row": _us_per_row(self.decode_duration_ms, self.decode_row_count), "write_duration_ms": round(self.write_duration_ms, 3), "write_row_count": self.write_row_count, "write_sample_count": self.write_sample_count, "write_us_per_row": _us_per_row(self.write_duration_ms, self.write_row_count), } def _us_per_row(duration_ms: float, row_count: int) -> float: if row_count <= 0: return 0.0 return round(duration_ms * 1000 / row_count, 3) def _sampled_stage_duration_ms(*, sample_duration_ms: float, sample_count: int, row_count: int) -> float: return round( _sampled_duration_ms(sample_duration_ms=sample_duration_ms, sample_count=sample_count, row_count=row_count), 3 ) def _sampled_duration_ms(*, sample_duration_ms: float, sample_count: int, row_count: int) -> float: if sample_count <= 0 or row_count <= 0: return 0.0 return sample_duration_ms * row_count / sample_count def _bounded_stage_durations_ms( *, total_duration_ms: float, decode_duration_ms: float, write_duration_ms: float, ) -> tuple[float, float, float]: total_duration_ms = max(total_duration_ms, 0.0) decode_duration_ms = min(max(decode_duration_ms, 0.0), total_duration_ms) remaining_duration_ms = max(total_duration_ms - decode_duration_ms, 0.0) write_duration_ms = min(max(write_duration_ms, 0.0), remaining_duration_ms) read_duration_ms = max(total_duration_ms - decode_duration_ms - write_duration_ms, 0.0) return round(read_duration_ms, 3), round(decode_duration_ms, 3), round(write_duration_ms, 3) def _operation_log_fields(stage_stats: SnapshotStageStats) -> dict[str, float | int]: return { "operation_upsert_count": stage_stats.write_row_count, "operation_upsert_duration_ms": round(stage_stats.write_duration_ms, 3), "operation_upsert_sample_count": stage_stats.write_sample_count, "operation_upsert_us_per_row": _us_per_row(stage_stats.write_duration_ms, stage_stats.write_row_count), } @dataclass(frozen=True) class SnapshotFileResult: rows_processed: int skipped_resume_rows: int duration_ms: float stage_stats: SnapshotStageStats stats: OperationStats @dataclass class SnapshotProgress: rows_processed: int = 0 files_processed: int = 0 table_count: int = 0 stage_stats: SnapshotStageStats = field(default_factory=SnapshotStageStats) def record_file(self, result: SnapshotFileResult) -> None: self.rows_processed += result.rows_processed self.files_processed += 1 self.stage_stats.merge(result.stage_stats) def _snapshot_row_header(first_row: list[str], expected_header: list[str]) -> tuple[list[str], bool]: if first_row == expected_header: return first_row, True if len(first_row) == len(expected_header) and set(first_row) == set(expected_header): raise SnapshotError("snapshot CSV header order does not match the SQL schema column order") return expected_header, False def _cursor_file(state: ConnectorState) -> str: return str(state.snapshot.get(SNAPSHOT_CURRENT_FILE) or "") def _cursor_next_row(state: ConnectorState) -> int: return int(state.snapshot.get(SNAPSHOT_NEXT_ROW) or 0) def _set_cursor( *, state: ConnectorState, file_key: str | None, next_row: int, ) -> None: if file_key: state.snapshot[SNAPSHOT_CURRENT_FILE] = file_key state.snapshot[SNAPSHOT_NEXT_ROW] = next_row return state.snapshot.pop(SNAPSHOT_CURRENT_FILE, None) state.snapshot.pop(SNAPSHOT_NEXT_ROW, None) def _resync_cursor_file(state: ConnectorState) -> str: resync = state.resync return str(resync.current_file or "") if resync is not None else "" def _resync_cursor_next_row(state: ConnectorState) -> int: resync = state.resync return int(resync.next_row or 0) if resync is not None else 0 def _set_resync_cursor( *, state: ConnectorState, file_key: str | None, next_row: int, ) -> None: resync = state.resync if resync is None: return resync.current_file = file_key resync.next_row = next_row def _set_snapshot_phase_cursor( *, state: ConnectorState, file_key: str | None, next_row: int, resync_snapshot_phase: bool, ) -> None: if resync_snapshot_phase: _set_resync_cursor(state=state, file_key=file_key, next_row=next_row) return _set_cursor(state=state, file_key=file_key, next_row=next_row) def _add_snapshot_diagnostic_columns(data: dict[str, Any], table: SnapshotTable) -> None: data["_incremental_commit_ts"] = 0 data["_incremental_source_schema"] = table.source_schema data["_incremental_source_table"] = table.source_table def _new_sink( *, config: ConnectorConfig, operations: OperationSink, table: SnapshotTable, row_decoder: CSVRowDecoder, ) -> SampledSink[Sequence[str | None]]: if config.operation_mode == OPERATION_MODE_STANDARD: data: dict[str, Any] = {} def decode(raw_values: Sequence[str | None]) -> dict[str, Any]: row_decoder.decode_into(raw_values, data) if config.enable_diagnostic_columns: _add_snapshot_diagnostic_columns(data, table) return data return DecodedSink( decode=decode, write_decoded=lambda decoded: operations.upsert(table=table.fivetran_table, data=decoded), ) return RawSink( write_raw=build_raw_csv_operation_writers( writer_factory=raw_csv_record_writer_factory(operations), table=table.fivetran_table, columns=table.columns, null_marker=config.snapshot.csv_null_value, binary_encoding=config.snapshot.binary_encoding, diagnostic_columns=( DiagnosticColumnSource( commit_ts=0, source_schema=table.source_schema, source_table=table.source_table, ) if config.enable_diagnostic_columns else None ), ).write_upsert ) def _sync_snapshot_file( *, config: ConnectorConfig, store: ObjectStore, state: ConnectorState, operations: OperationSink, policy: CheckpointPolicy, table: SnapshotTable, sink: SampledSink[Sequence[str | None]], file_key: str, cursor_file: str, cursor_next_row: int, resync_snapshot_phase: bool, ) -> SnapshotFileResult: file_start = timer_start() file_stats = OperationStats() file_operations = TimedOperationSink(operations, file_stats) rows_processed = 0 skipped_resume_rows = 0 stage_stats = SnapshotStageStats() next_row = cursor_next_row if file_key == cursor_file else 0 _set_snapshot_phase_cursor( state=state, file_key=file_key, next_row=next_row, resync_snapshot_phase=resync_snapshot_phase, ) expected_header = table.column_names with open_csv_text( store, file_key, buffer_s3=config.snapshot.s3_csv_read_mode == S3_CSV_READ_MODE_BUFFERED, ) as handle: reader = iter_csv_rows( handle, parser=new_csv_parser(error_type=SnapshotError, context="snapshot", strict=True), escape_backslash=config.snapshot.escape_backslash, ) first_row = next(reader, None) if first_row is None: rows = iter(()) else: _, has_header = _snapshot_row_header(first_row, expected_header) rows = reader if has_header else chain([first_row], reader) decode_sample_duration_ms = 0.0 write_sample_duration_ms = 0.0 sample_count = 0 next_sample_row = 0 read_row_count = 0 for row_offset, raw_values in enumerate(rows): read_row_count += 1 if row_offset < next_row: skipped_resume_rows += 1 continue if rows_processed == next_sample_row: sample = sink.write_sample(raw_values) decode_sample_duration_ms += sample.decode_duration_ms write_sample_duration_ms += sample.write_duration_ms sample_count += 1 next_sample_row += _SNAPSHOT_STAGE_SAMPLE_INTERVAL_ROWS else: sink.write(raw_values) rows_processed += 1 consumed_rows = skipped_resume_rows + rows_processed data_loop_duration_ms = elapsed_ms(file_start) stage_stats.read_row_count = read_row_count stage_stats.decode_row_count = rows_processed stage_stats.write_row_count = rows_processed stage_stats.decode_sample_count = sample_count stage_stats.write_sample_count = sample_count decode_duration_ms = _sampled_stage_duration_ms( sample_duration_ms=decode_sample_duration_ms, sample_count=sample_count, row_count=rows_processed, ) write_duration_ms = _sampled_stage_duration_ms( sample_duration_ms=write_sample_duration_ms, sample_count=sample_count, row_count=rows_processed, ) stage_stats.read_duration_ms, stage_stats.decode_duration_ms, stage_stats.write_duration_ms = ( _bounded_stage_durations_ms( total_duration_ms=data_loop_duration_ms, decode_duration_ms=decode_duration_ms, write_duration_ms=write_duration_ms, ) ) _set_snapshot_phase_cursor( state=state, file_key=file_key, next_row=consumed_rows, resync_snapshot_phase=resync_snapshot_phase, ) file_operations.checkpoint(state=state) policy.mark_checkpointed() return SnapshotFileResult( rows_processed=rows_processed, skipped_resume_rows=skipped_resume_rows, duration_ms=elapsed_ms(file_start), stage_stats=stage_stats, stats=file_stats, ) def _table_name_messages(table_names: set[TableName]) -> list[str]: return sorted(table_name.quoted for table_name in table_names) def _table_in_scope(table: SnapshotTable, resync_tables: set[TableName]) -> bool: return not resync_tables or table.name in resync_tables def _truncate_snapshot_targets( *, store: ObjectStore, snapshot_prefix: str, resync_tables: set[TableName], operations: TimedOperationSink, ) -> int: truncated = 0 found_resync_tables: set[TableName] = set() for table in iter_snapshot_tables(store, snapshot_prefix): if not _table_in_scope(table, resync_tables): continue found_resync_tables.add(table.name) unsupported_primary_key_columns = _unsupported_primary_key_columns(table) if unsupported_primary_key_columns: warning( "skipping snapshot table because primary key contains unsupported column type", phase="snapshot", source_schema=table.source_schema, source_table=table.source_table, columns=unsupported_primary_key_columns, ) continue operations.truncate(table=table.fivetran_table) truncated += 1 missing_tables = resync_tables - found_resync_tables if missing_tables: raise SnapshotError(f"resync snapshot prefix is missing target tables: {_table_name_messages(missing_tables)}") return truncated def _log_snapshot_file_completed(table: SnapshotTable, file_key: str, result: SnapshotFileResult) -> None: info_with_memory( "snapshot file completed", phase="snapshot", source_schema=table.source_schema, source_table=table.source_table, fivetran_table=table.fivetran_table, file_path=file_key, duration_ms=result.duration_ms, rows_processed=result.rows_processed, skipped_resume_rows=result.skipped_resume_rows, **result.stage_stats.as_log_fields(), **_operation_log_fields(result.stage_stats), **result.stats.as_log_fields(prefix="timed_operation_"), ) def _sync_snapshot_tables( *, config: ConnectorConfig, store: ObjectStore, state: ConnectorState, operations: OperationSink, policy: CheckpointPolicy, snapshot_prefix: str, resync_tables: set[TableName], cursor_file: str, cursor_next_row: int, resync_snapshot_phase: bool, phase_stats: OperationStats, progress: SnapshotProgress, ) -> None: found_resync_tables: set[TableName] = set() cursor_found = not cursor_file for table in iter_snapshot_tables(store, snapshot_prefix): if resync_tables and table.name not in resync_tables: continue found_resync_tables.add(table.name) unsupported_primary_key_columns = _unsupported_primary_key_columns(table) if unsupported_primary_key_columns: warning( "skipping snapshot table because primary key contains unsupported column type", phase="snapshot", source_schema=table.source_schema, source_table=table.source_table, columns=unsupported_primary_key_columns, ) continue progress.table_count += 1 row_decoder = CSVRowDecoder.new( table.columns, null_marker=config.snapshot.csv_null_value, binary_encoding=config.snapshot.binary_encoding, ) sink = _new_sink( config=config, operations=operations, table=table, row_decoder=row_decoder, ) for file_key in _snapshot_files(store, table, snapshot_prefix): if not cursor_found: if file_key != cursor_file: continue cursor_found = True file_result = _sync_snapshot_file( config=config, store=store, state=state, operations=operations, policy=policy, table=table, sink=sink, file_key=file_key, cursor_file=cursor_file, cursor_next_row=cursor_next_row, resync_snapshot_phase=resync_snapshot_phase, ) phase_stats.merge(file_result.stats) progress.record_file(file_result) _log_snapshot_file_completed(table, file_key, file_result) if not cursor_found: raise SnapshotError(f"snapshot cursor file {cursor_file!r} is not present") missing_resync_tables = resync_tables - found_resync_tables if missing_resync_tables: raise SnapshotError( f"resync snapshot prefix is missing syncable target tables: {_table_name_messages(missing_resync_tables)}" ) if progress.table_count == 0: raise SnapshotError("no supported tables discovered from snapshot schema files") def _finish_snapshot_phase( *, state: ConnectorState, operations: OperationSink, policy: CheckpointPolicy, resync_tables: set[TableName], snapshot_tso: int, ) -> None: if resync_tables: _set_resync_cursor(state=state, file_key=None, next_row=0) if finish_table_resync_if_complete(state): operations.checkpoint(state=state) policy.mark_checkpointed() return state.mark_snapshot_done() if state.checkpoint_ts < snapshot_tso: state.checkpoint_ts = snapshot_tso _set_cursor(state=state, file_key=None, next_row=0) operations.checkpoint(state=state) policy.mark_checkpointed() def sync_snapshot( *, config: ConnectorConfig, store: ObjectStore, state: ConnectorState, operations: OperationSink, ) -> None: phase_start = timer_start() phase_stats = OperationStats() phase_operations = TimedOperationSink(operations, phase_stats) progress = SnapshotProgress() resync = state.resync resync_snapshot = resync if resync is not None and resync.is_snapshot_phase() else None resync_snapshot_phase = resync_snapshot is not None if resync_snapshot is None and snapshots_complete(state): return snapshot_prefix = ( str(resync_snapshot.snapshot_prefix).strip("/") if resync_snapshot is not None else config.snapshot_prefix ) if not store.prefix_exists(snapshot_prefix): raise SnapshotError(f"snapshot prefix does not exist: {snapshot_prefix}") snapshot_tso = snapshot_metadata_pos(store, snapshot_prefix) if resync_snapshot is not None and snapshot_tso != resync_snapshot.snapshot_tso: warning( "resync snapshot_tso differs from snapshot metadata; using metadata Pos as ground truth", phase="snapshot", request_id=resync_snapshot.request_id, state_snapshot_tso=resync_snapshot.snapshot_tso, metadata_snapshot_tso=snapshot_tso, ) resync_snapshot.snapshot_tso = snapshot_tso phase_operations.checkpoint(state=state) if resync_snapshot is None and state_snapshot_tso(state) <= 0 and set_snapshot_tso(state, snapshot_tso): phase_operations.checkpoint(state=state) policy = CheckpointPolicy(config.checkpoint.interval_rows, config.checkpoint.interval_seconds) resync_tables = table_resync_names(state) cursor_file = _resync_cursor_file(state) if resync_snapshot_phase else _cursor_file(state) cursor_next_row = _resync_cursor_next_row(state) if resync_snapshot_phase else _cursor_next_row(state) info_with_memory( "starting snapshot phase", phase="snapshot", snapshot_tso=snapshot_tso, snapshot_prefix=snapshot_prefix, ) try: if not resync_snapshot_phase and config.truncate_before_snapshot and not state.snapshot_truncate_done: truncated = _truncate_snapshot_targets( store=store, snapshot_prefix=snapshot_prefix, resync_tables=set(), operations=phase_operations, ) state.mark_snapshot_truncated() phase_operations.checkpoint(state=state) policy.mark_checkpointed() info( "snapshot target tables truncated before full snapshot", phase="snapshot", table_count=truncated, ) if resync_tables and not table_resync_truncate_done(state): truncated = _truncate_snapshot_targets( store=store, snapshot_prefix=snapshot_prefix, resync_tables=resync_tables, operations=phase_operations, ) if mark_table_resync_truncated(state): phase_operations.checkpoint(state=state) policy.mark_checkpointed() info( "resync target tables truncated before snapshot", phase="snapshot", table_count=truncated, ) _sync_snapshot_tables( config=config, store=store, state=state, operations=operations, policy=policy, snapshot_prefix=snapshot_prefix, resync_tables=resync_tables, cursor_file=cursor_file, cursor_next_row=cursor_next_row, resync_snapshot_phase=resync_snapshot_phase, phase_stats=phase_stats, progress=progress, ) _finish_snapshot_phase( state=state, operations=phase_operations, policy=policy, resync_tables=resync_tables, snapshot_tso=snapshot_tso, ) except Exception as exc: warning_with_memory( "snapshot phase failed", phase="snapshot", duration_ms=elapsed_ms(phase_start), rows_processed=progress.rows_processed, files_processed=progress.files_processed, error_type=type(exc).__name__, error_message=str(exc), **progress.stage_stats.as_log_fields(), **_operation_log_fields(progress.stage_stats), **phase_stats.as_log_fields(prefix="timed_operation_"), ) raise phase_duration_ms = elapsed_ms(phase_start) info_with_memory( "snapshot phase completed", phase="snapshot", snapshot_tso=snapshot_tso, duration_ms=phase_duration_ms, rows_processed=progress.rows_processed, files_processed=progress.files_processed, table_count=progress.table_count, **progress.stage_stats.as_log_fields(), **_operation_log_fields(progress.stage_stats), **phase_stats.as_log_fields(prefix="timed_operation_"), ) __all__ = ["sync_snapshot"]