def update_bookmark(row, state, tap_stream_id, replication_key_name):
replication_key_value = row.get(replication_key_name)
if replication_key_value:
replication_key_type = replication_key_value.__class__.__name__
replication_key_value_bookmark = common.class_to_string(
replication_key_value, replication_key_type)
state = singer.write_bookmark(state, tap_stream_id,
'replication_key_value',
replication_key_value_bookmark)
state = singer.write_bookmark(state, tap_stream_id,
'replication_key_type',
replication_key_type)
def update_bookmark(row: Dict, state: Dict, tap_stream_id: str,
replication_key_name: str) -> None:
"""
Updates replication key and type values in state bookmark
Args:
row: DB record
state: dictionary of bookmarks
tap_stream_id: stream ID
replication_key_name: replication key
"""
replication_key_value = row.get(replication_key_name)
if replication_key_value:
replication_key_type = replication_key_value.__class__.__name__
replication_key_value_bookmark = common.class_to_string(
replication_key_value, replication_key_type)
state = singer.write_bookmark(state, tap_stream_id,
'replication_key_value',
replication_key_value_bookmark)
singer.write_bookmark(state, tap_stream_id, 'replication_key_type',
replication_key_type)
def sync_collection(client, stream, state, projection):
tap_stream_id = stream['tap_stream_id']
LOGGER.info('Starting full table sync for %s', tap_stream_id)
md_map = metadata.to_map(stream['metadata'])
database_name = metadata.get(md_map, (), 'database-name')
db = client[database_name]
collection = db[stream['stream']]
#before writing the table version to state, check if we had one to begin with
first_run = singer.get_bookmark(state, stream['tap_stream_id'],
'version') is None
# last run was interrupted if there is a last_id_fetched bookmark
was_interrupted = singer.get_bookmark(state, stream['tap_stream_id'],
'last_id_fetched') is not None
#pick a new table version if last run wasn't interrupted
if was_interrupted:
stream_version = singer.get_bookmark(state, stream['tap_stream_id'],
'version')
else:
stream_version = int(time.time() * 1000)
state = singer.write_bookmark(state, stream['tap_stream_id'], 'version',
stream_version)
singer.write_message(singer.StateMessage(value=copy.deepcopy(state)))
activate_version_message = singer.ActivateVersionMessage(
stream=common.calculate_destination_stream_name(stream),
version=stream_version)
# For the initial replication, emit an ACTIVATE_VERSION message
# at the beginning so the records show up right away.
if first_run:
singer.write_message(activate_version_message)
if singer.get_bookmark(state, stream['tap_stream_id'], 'max_id_value'):
# There is a bookmark
max_id_value = singer.get_bookmark(state, stream['tap_stream_id'],
'max_id_value')
max_id_type = singer.get_bookmark(state, stream['tap_stream_id'],
'max_id_type')
max_id_value = common.string_to_class(max_id_value, max_id_type)
else:
max_id_value = get_max_id_value(collection)
last_id_fetched = singer.get_bookmark(state, stream['tap_stream_id'],
'last_id_fetched')
if max_id_value:
# Write the bookmark if max_id_value is defined
state = singer.write_bookmark(
state, stream['tap_stream_id'], 'max_id_value',
common.class_to_string(max_id_value,
max_id_value.__class__.__name__))
state = singer.write_bookmark(state, stream['tap_stream_id'],
'max_id_type',
max_id_value.__class__.__name__)
find_filter = {'$lte': max_id_value}
if last_id_fetched:
last_id_fetched_type = singer.get_bookmark(state,
stream['tap_stream_id'],
'last_id_fetched_type')
find_filter['$gte'] = common.string_to_class(last_id_fetched,
last_id_fetched_type)
query_message = 'Querying {} with:\n\tFind Parameters: {}'.format(
stream['tap_stream_id'], find_filter)
if projection:
query_message += '\n\tProjection: {}'.format(projection)
# pylint: disable=logging-format-interpolation
LOGGER.info(query_message)
with collection.find({'_id': find_filter},
projection,
sort=[("_id", pymongo.ASCENDING)]) as cursor:
rows_saved = 0
time_extracted = utils.now()
start_time = time.time()
schema = stream['schema'] or {"type": "object", "properties": {}}
for row in cursor:
rows_saved += 1
schema_build_start_time = time.time()
if common.row_to_schema(schema, row):
singer.write_message(
singer.SchemaMessage(
stream=common.calculate_destination_stream_name(
stream),
schema=schema,
key_properties=['_id']))
common.SCHEMA_COUNT[stream['tap_stream_id']] += 1
common.SCHEMA_TIMES[stream['tap_stream_id']] += time.time(
) - schema_build_start_time
record_message = common.row_to_singer_record(
stream, row, stream_version, time_extracted)
singer.write_message(record_message)
state = singer.write_bookmark(
state, stream['tap_stream_id'], 'last_id_fetched',
common.class_to_string(row['_id'],
row['_id'].__class__.__name__))
state = singer.write_bookmark(state, stream['tap_stream_id'],
'last_id_fetched_type',
row['_id'].__class__.__name__)
if rows_saved % common.UPDATE_BOOKMARK_PERIOD == 0:
singer.write_message(
singer.StateMessage(value=copy.deepcopy(state)))
common.COUNTS[tap_stream_id] += rows_saved
common.TIMES[tap_stream_id] += time.time() - start_time
# clear max pk value and last pk fetched upon successful sync
singer.clear_bookmark(state, stream['tap_stream_id'], 'max_id_value')
singer.clear_bookmark(state, stream['tap_stream_id'], 'max_id_type')
singer.clear_bookmark(state, stream['tap_stream_id'], 'last_id_fetched')
singer.clear_bookmark(state, stream['tap_stream_id'],
'last_id_fetched_type')
state = singer.write_bookmark(state, stream['tap_stream_id'],
'initial_full_table_complete', True)
singer.write_message(activate_version_message)
LOGGER.info('Syncd {} records for {}'.format(rows_saved, tap_stream_id))
def test_class_to_string_with_unsupported_type_raises_exception(self):
with self.assertRaises(UnsupportedKeyTypeException):
common.class_to_string('a', 'random type')
def test_class_to_string_with_bson_Timestamp_should_return_concatenated_time(
self):
ts = bson.Timestamp(200000, 80)
self.assertEqual('200000.80', common.class_to_string(ts, 'Timestamp'))
def sync_collection(collection: Collection, stream: Dict, state: Dict) -> None:
"""
Sync collection records incrementally
Args:
collection: MongoDB collection instance
stream: dictionary of all stream details
state: the tap state
"""
LOGGER.info('Starting full table sync for %s', stream['tap_stream_id'])
# before writing the table version to state, check if we had one to begin with
first_run = singer.get_bookmark(state, stream['tap_stream_id'],
'version') is None
# last run was interrupted if there is a last_id_fetched bookmark
# pick a new table version if last run wasn't interrupted
if singer.get_bookmark(state, stream['tap_stream_id'],
'last_id_fetched') is not None:
stream_version = singer.get_bookmark(state, stream['tap_stream_id'],
'version')
else:
stream_version = int(time.time() * 1000)
state = singer.write_bookmark(state, stream['tap_stream_id'], 'version',
stream_version)
singer.write_message(singer.StateMessage(value=copy.deepcopy(state)))
activate_version_message = singer.ActivateVersionMessage(
stream=common.calculate_destination_stream_name(stream),
version=stream_version)
# For the initial replication, emit an ACTIVATE_VERSION message
# at the beginning so the records show up right away.
if first_run:
singer.write_message(activate_version_message)
if singer.get_bookmark(state, stream['tap_stream_id'], 'max_id_value'):
# There is a bookmark
max_id_value = common.string_to_class(
singer.get_bookmark(state, stream['tap_stream_id'],
'max_id_value'),
singer.get_bookmark(state, stream['tap_stream_id'], 'max_id_type'))
else:
max_id_value = get_max_id_value(collection)
last_id_fetched = singer.get_bookmark(state, stream['tap_stream_id'],
'last_id_fetched')
if max_id_value:
# Write the bookmark if max_id_value is defined
state = singer.write_bookmark(
state, stream['tap_stream_id'], 'max_id_value',
common.class_to_string(max_id_value,
max_id_value.__class__.__name__))
state = singer.write_bookmark(state, stream['tap_stream_id'],
'max_id_type',
max_id_value.__class__.__name__)
find_filter = {'$lte': max_id_value}
if last_id_fetched:
find_filter['$gte'] = common.string_to_class(
last_id_fetched,
singer.get_bookmark(state, stream['tap_stream_id'],
'last_id_fetched_type'))
LOGGER.info('Querying %s with: %s', stream['tap_stream_id'],
dict(find=find_filter))
with collection.find({'_id': find_filter},
sort=[("_id", pymongo.ASCENDING)]) as cursor:
rows_saved = 0
start_time = time.time()
for row in cursor:
rows_saved += 1
singer.write_message(
common.row_to_singer_record(stream=stream,
row=row,
time_extracted=utils.now(),
time_deleted=None,
version=stream_version))
state = singer.write_bookmark(
state, stream['tap_stream_id'], 'last_id_fetched',
common.class_to_string(row['_id'],
row['_id'].__class__.__name__))
state = singer.write_bookmark(state, stream['tap_stream_id'],
'last_id_fetched_type',
row['_id'].__class__.__name__)
if rows_saved % common.UPDATE_BOOKMARK_PERIOD == 0:
singer.write_message(
singer.StateMessage(value=copy.deepcopy(state)))
common.COUNTS[stream['tap_stream_id']] += rows_saved
common.TIMES[stream['tap_stream_id']] += time.time() - start_time
# clear max pk value and last pk fetched upon successful sync
singer.clear_bookmark(state, stream['tap_stream_id'], 'max_id_value')
singer.clear_bookmark(state, stream['tap_stream_id'], 'max_id_type')
singer.clear_bookmark(state, stream['tap_stream_id'], 'last_id_fetched')
singer.clear_bookmark(state, stream['tap_stream_id'],
'last_id_fetched_type')
singer.write_bookmark(state, stream['tap_stream_id'],
'initial_full_table_complete', True)
singer.write_message(activate_version_message)
LOGGER.info('Syncd %s records for %s', rows_saved, stream['tap_stream_id'])