def _get_data(query):
"""
Execute the query and return the result as a list of dictionaries
:param query:
:return:
"""
handler = GOBStorageHandler()
with handler.get_session() as session:
data = session.execute(query)
for row in data:
yield _convert_row(row)
def apply(msg):
mode = msg['header'].get('mode', FULL_UPLOAD)
logger.configure(msg, "UPDATE")
logger.info("Apply events")
storage = GOBStorageHandler()
combinations = _get_source_catalog_entity_combinations(storage, msg)
# Gather statistics of update process
stats = UpdateStatistics()
before = None
after = None
for result in combinations:
model = f"{result.source} {result.catalogue} {result.entity}"
logger.info(f"Apply events {model}")
storage = GOBStorageHandler(result)
# Track eventId before event application
entity_max_eventid, last_eventid = get_event_ids(storage)
before = min(entity_max_eventid or 0, before or sys.maxsize)
if is_corrupted(entity_max_eventid, last_eventid):
logger.error(
f"Model {model} is inconsistent! data is more recent than events"
)
elif entity_max_eventid == last_eventid:
logger.info(f"Model {model} is up to date")
apply_confirm_events(storage, stats, msg)
else:
logger.info(f"Start application of unhandled {model} events")
with storage.get_session():
last_events = storage.get_last_events(
) # { tid: last_event, ... }
apply_events(storage, last_events, entity_max_eventid, stats)
apply_confirm_events(storage, stats, msg)
# Track eventId after event application
entity_max_eventid, last_eventid = get_event_ids(storage)
after = max(entity_max_eventid or 0, after or 0)
# Build result message
results = stats.results()
if mode == FULL_UPLOAD and _should_analyze(stats):
logger.info("Running VACUUM ANALYZE on table")
storage.analyze_table()
stats.log()
logger.info(f"Apply events {model} completed", {'data': results})
msg['summary'] = logger.get_summary()
# Add a events notification telling what types of event have been applied
if not msg['header'].get('suppress_notifications', False):
add_notification(msg, EventNotification(stats.applied,
[before, after]))
return msg
def _execute_multiple(queries, stream=False, max_row_buffer=1000):
handler = GOBStorageHandler()
with handler.get_session() as session:
if stream:
connection = session.connection(execution_options={
'stream_results': True,
'max_row_buffer': max_row_buffer
})
execute_on = connection
else:
execute_on = session
# Commit all queries as a whole on exit with
for query in queries:
result = execute_on.execute(query)
return result # Return result of last execution
def setUp(self):
self.mock_model = MagicMock(spec=GOBModel)
self.msg = fixtures.get_message_fixture()
model = {
"entity_id": "identificatie",
"version": "1",
"has_states": False,
}
# Add the hash to the message
populator = Populator(model, self.msg)
self.msg['header']['source'] = 'any source'
for content in self.msg['contents']:
populator.populate(content)
message = ImportMessage(self.msg)
metadata = message.metadata
GOBStorageHandler.base = MagicMock()
self.storage = GOBStorageHandler(metadata)
GOBStorageHandler.engine = MagicMock()
GOBStorageHandler.engine.__enter__ = lambda self: None
GOBStorageHandler.engine.__exit__ = lambda *args: None
GOBStorageHandler.engine.begin = lambda: GOBStorageHandler.engine
def _find_occurrence_of_column(catalog_name: str, collection_name: str, column: str, first_or_last: str,
storage: GOBStorageHandler):
assert first_or_last in ('first', 'last')
query = f"""
SELECT eventid FROM events
WHERE catalogue='{catalog_name}' AND entity='{collection_name}' AND (
(action='ADD' AND (contents->'entity')::jsonb ? '{column}')
OR
(action='MODIFY' AND (contents->>'modifications')::jsonb @> '[{{\"key\": \"{column}\"}}]')
)
ORDER BY eventid {'ASC' if first_or_last == 'first' else 'DESC'}
LIMIT 1
"""
with storage.get_session() as session:
try:
return next(session.execute(query))[0]
except StopIteration:
return None
def full_update(msg):
"""Store the events for the current dataset
:param msg: the result of the application of the events
:return: Result message
"""
logger.configure(msg, "UPDATE")
logger.info(
f"Update to GOB Database {GOBStorageHandler.user_name} started")
# Interpret the message header
message = ImportMessage(msg)
metadata = message.metadata
storage = GOBStorageHandler(metadata)
model = f"{metadata.source} {metadata.catalogue} {metadata.entity}"
logger.info(f"Store events {model}")
# Get events from message
events = msg["contents"]
# Gather statistics of update process
stats = UpdateStatistics()
_process_events(storage, events, stats)
# Build result message
results = stats.results()
stats.log()
logger.info(f"Store events {model} completed", {'data': results})
results.update(logger.get_summary())
# Return the result message, with no log, no contents but pass-through any confirms
message = {
"header": msg["header"],
"summary": results,
"contents": None,
"confirms": msg.get('confirms')
}
return message
def correct_version_numbers():
storage = GOBStorageHandler()
for catalog_name, collection in GOBMigrations()._migrations.items():
for collection_name, versions in collection.items():
print(f"{catalog_name} {collection_name}: Determine version boundaries in events")
current_version = '0.1'
change_eventids = [(current_version, 0)]
migration = versions[current_version]
while migration:
change_eventid = _get_change_eventid(catalog_name, collection_name, migration, storage)
target_version = migration['target_version']
if change_eventid is not None:
change_eventids.append((target_version, change_eventid))
migration = versions.get(target_version)
_update_version_numbers(catalog_name, collection_name, change_eventids, storage)
def update_materialized_view(msg):
"""Updates materialized view for a relation for a given catalog, collection and attribute or relation name.
Expects a message with headers:
- catalogue
- collection (if catalogue is 'rel' this should be the relation_name)
- attribute (optional if catalogue is 'rel')
examples of correct headers that are functionally equivalent:
header = {
"catalogue": "meetbouten",
"collection": "meetbouten",
"attribute": "ligt_in_buurt",
}
header = {
"catalogue": "rel",
"collection": "mbn_mbt_gbd_brt_ligt_in_buurt",
}
:param msg:
:return:
"""
header = msg.get('header', {})
catalog_name = header.get('catalogue')
collection_name = header.get('collection')
attribute_name = header.get('attribute')
logger.configure(msg, "UPDATE_VIEW")
storage_handler = GOBStorageHandler()
view = _get_materialized_view(catalog_name, collection_name, attribute_name)
view.refresh(storage_handler)
logger.info(f"Update materialized view {view.name}")
timestamp = datetime.datetime.utcnow().isoformat()
msg['header'].update({
"timestamp": timestamp
})
return msg
def compare(msg):
"""Compare new data in msg (contents) with the current data
:param msg: The new data, including header and summary
:return: result message
"""
logger.configure(msg, "COMPARE")
header = msg.get('header', {})
mode = header.get('mode', FULL_UPLOAD)
logger.info(
f"Compare (mode = {mode}) to GOB Database {GOBStorageHandler.user_name} started"
)
# Parse the message header
message = ImportMessage(msg)
metadata = message.metadata
# Get the model for the collection to be compared
gob_model = GOBModel()
entity_model = gob_model.get_collection(metadata.catalogue,
metadata.entity)
# Initialize a storage handler for the collection
storage = GOBStorageHandler(metadata)
model = f"{metadata.source} {metadata.catalogue} {metadata.entity}"
logger.info(f"Compare {model}")
stats = CompareStatistics()
tmp_table_name = None
with storage.get_session():
with ProgressTicker("Collect compare events", 10000) as progress:
# Check any dependencies
if not meets_dependencies(storage, msg):
return {
"header": msg["header"],
"summary": logger.get_summary(),
"contents": None
}
enricher = Enricher(storage, msg)
populator = Populator(entity_model, msg)
# If there are no records in the database all data are ADD events
initial_add = not storage.has_any_entity()
if initial_add:
logger.info("Initial load of new collection detected")
# Write ADD events directly, without using a temporary table
contents_writer = ContentsWriter()
contents_writer.open()
# Pass a None confirms_writer because only ADD events are written
collector = EventCollector(contents_writer,
confirms_writer=None,
version=entity_model['version'])
collect = collector.collect_initial_add
else:
# Collect entities in a temporary table
collector = EntityCollector(storage)
collect = collector.collect
tmp_table_name = collector.tmp_table_name
for entity in msg["contents"]:
progress.tick()
stats.collect(entity)
enricher.enrich(entity)
populator.populate(entity)
collect(entity)
collector.close()
if initial_add:
filename = contents_writer.filename
confirms = None
contents_writer.close()
else:
# Compare entities from temporary table
with storage.get_session():
diff = storage.compare_temporary_data(tmp_table_name, mode)
filename, confirms = _process_compare_results(
storage, entity_model, diff, stats)
# Build result message
results = stats.results()
logger.info(f"Compare {model} completed", {'data': results})
results.update(logger.get_summary())
message = {
"header": msg["header"],
"summary": results,
"contents_ref": filename,
"confirms": confirms
}
return message
default=False,
help='migrate the database tables, views and indexes')
parser.add_argument('--materialized_views',
action='store_true',
default=False,
help='force recreation of materialized views')
parser.add_argument('mv_name',
nargs='?',
help='The materialized view to update. Use with --materialized-views')
args = parser.parse_args()
if DEBUG:
print("WARNING: Debug mode is ON")
# Initialize database tables
storage = GOBStorageHandler()
# Migrate on request only
if args.migrate:
recreate = [args.mv_name] if args.materialized_views and args.mv_name else args.materialized_views
storage.init_storage(force_migrate=True, recreate_materialized_views=recreate)
else:
storage.init_storage()
params = {
"stream_contents": True,
"thread_per_service": True,
APPLY_QUEUE: {
"load_message": False
}
}
MessagedrivenService(SERVICEDEFINITION, "Upload", params).start()
def test_base(self):
GOBStorageHandler.base = None
GOBStorageHandler._set_base(update=True)
self.assertIsNotNone(GOBStorageHandler.base)
class TestStorageHandler(unittest.TestCase):
@patch('gobupload.storage.handler.create_engine', MagicMock())
def setUp(self):
self.mock_model = MagicMock(spec=GOBModel)
self.msg = fixtures.get_message_fixture()
model = {
"entity_id": "identificatie",
"version": "1",
"has_states": False,
}
# Add the hash to the message
populator = Populator(model, self.msg)
self.msg['header']['source'] = 'any source'
for content in self.msg['contents']:
populator.populate(content)
message = ImportMessage(self.msg)
metadata = message.metadata
GOBStorageHandler.base = MagicMock()
self.storage = GOBStorageHandler(metadata)
GOBStorageHandler.engine = MagicMock()
GOBStorageHandler.engine.__enter__ = lambda self: None
GOBStorageHandler.engine.__exit__ = lambda *args: None
GOBStorageHandler.engine.begin = lambda: GOBStorageHandler.engine
@patch("gobupload.storage.handler.automap_base", MagicMock())
def test_base(self):
GOBStorageHandler.base = None
GOBStorageHandler._set_base(update=True)
self.assertIsNotNone(GOBStorageHandler.base)
@patch("gobupload.storage.handler.alembic.config")
@patch('gobupload.storage.handler.alembic.script')
@patch("gobupload.storage.handler.migration")
def test_init_storage(self, mock_alembic, mock_script, mock_config):
context = MagicMock()
context.get_current_revision.return_value = "revision 1"
mock_alembic.MigrationContext.configure.return_value = context
script = MagicMock()
script.get_current_head.return_value = "revision 2"
mock_script.ScriptDirectory.from_config.return_value = script
self.storage._init_views = MagicMock()
self.storage._get_reflected_base = MagicMock()
self.storage._init_indexes = MagicMock()
self.storage._set_base = MagicMock()
self.storage._init_relation_materialized_views = MagicMock()
self.storage._check_configuration = MagicMock()
self.storage.init_storage(recreate_materialized_views='booleanValue')
# mock_alembic.config.main.assert_called_once()
self.storage._init_views.assert_called_once()
# self.storage._set_base.assert_called_with(update=True)
self.storage._init_indexes.assert_called_once()
self.storage._init_relation_materialized_views.assert_called_with(
'booleanValue')
self.storage._check_configuration.assert_called_once()
def test_get_config_value(self):
self.storage.engine = MagicMock()
self.storage.engine.execute.return_value = iter([('the value', )])
self.assertEqual('the value',
self.storage._get_config_value('the setting'))
self.storage.engine.execute.assert_called_with('SHOW the setting')
@patch("builtins.print")
def test_check_configuration(self, mock_print):
self.storage._get_config_value = lambda x: 'the value'
self.storage.config_checks = [('the setting', lambda x: True,
'the message', self.storage.WARNING)]
self.storage._check_configuration()
mock_print.assert_not_called()
self.storage.config_checks = [('the setting', lambda x: False,
'the message', self.storage.WARNING)]
self.storage._check_configuration()
mock_print.assert_called_with(
'WARNING: Checking Postgres config for the setting. '
'Value is the value, but the message')
mock_print.reset_mock()
self.storage.config_checks = [('the setting', lambda x: False,
'the message', self.storage.ERROR)]
with self.assertRaises(GOBException):
self.storage._check_configuration()
@patch("gobupload.storage.handler.MaterializedViews")
def test_init_relation_materialized_view(self, mock_materialized_views):
self.storage._init_relation_materialized_views()
mock_materialized_views.assert_called_once()
mock_materialized_views.return_value.initialise.assert_called_with(
self.storage, False)
mock_materialized_views.reset_mock()
self.storage._init_relation_materialized_views(True)
mock_materialized_views.assert_called_once()
mock_materialized_views.return_value.initialise.assert_called_with(
self.storage, True)
def test_indexes_to_drop_query(self):
expected = """
SELECT
s.indexrelname
FROM pg_catalog.pg_stat_user_indexes s
JOIN pg_catalog.pg_index i ON s.indexrelid = i.indexrelid
WHERE
s.relname in ('sometable_1','sometable_2')
AND s.indexrelname not in ('index_a','index_b')
AND 0 <> ALL (i.indkey) -- no index column is an expression
AND NOT i.indisunique -- no unique indexes
AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_constraint c WHERE c.conindid = s.indexrelid)
"""
self.assertEqual(
expected,
self.storage._indexes_to_drop_query(['sometable_1', 'sometable_2'],
['index_a', 'index_b']))
def test_drop_indexes(self):
gobupload.storage.handler.indexes = {
"index_a": {
"table_name": "sometable_1",
"columns": ["col_a", "col_b"],
},
"index_b": {
"table_name": "sometable_2",
"columns": ["col_a", "col_b"],
},
}
self.storage.engine = MagicMock()
self.storage._indexes_to_drop_query = MagicMock()
self.storage.engine.execute.return_value = [("index_c", ),
("index_d", )]
self.storage.execute = MagicMock()
self.storage._drop_indexes()
self.storage.execute.assert_has_calls([
call('DROP INDEX IF EXISTS "index_c"'),
call('DROP INDEX IF EXISTS "index_d"'),
])
self.storage._indexes_to_drop_query.assert_called_with(
['sometable_1', 'sometable_2'],
['index_a', 'index_b'],
)
self.storage.engine.execute.assert_called_with(
self.storage._indexes_to_drop_query.return_value)
@patch("builtins.print")
def test_drop_indexes_exception_get(self, mock_print):
e = OperationalError('stmt', 'params', 'orig')
self.storage.engine.execute = MagicMock(side_effect=e)
self.storage._indexes_to_drop_query = MagicMock()
self.storage._drop_indexes()
mock_print.assert_called_with(
f"ERROR: Could not get indexes to drop: {str(e)}")
@patch("builtins.print")
def test_drop_indexes_exception_drop(self, mock_print):
e = OperationalError('stmt', 'params', 'orig')
self.storage.engine.execute = MagicMock(return_value=[('a', )])
self.storage._indexes_to_drop_query = MagicMock()
self.storage.execute = MagicMock(side_effect=e)
self.storage._drop_indexes()
mock_print.assert_called_with(
f"ERROR: Could not drop index a: {str(e)}")
def test_get_existing_indexes(self):
self.storage.engine.execute = MagicMock(
return_value=[('indexA', ), ('indexB', )])
self.assertEqual(['indexA', 'indexB'],
self.storage._get_existing_indexes())
@patch("builtins.print")
def test_get_existing_indexes_exception(self, mock_print):
e = OperationalError('stmt', 'params', 'orig')
self.storage.engine.execute = MagicMock(side_effect=e)
self.assertEqual([], self.storage._get_existing_indexes())
mock_print.assert_called_with(
f"WARNING: Could not fetch list of existing indexes: {e}")
def test_init_indexes(self):
gobupload.storage.handler.indexes = {
"indexname": {
"table_name": "sometable",
"columns": ["cola", "colb"],
},
"index2name": {
"table_name": "someothertable",
"columns": ["cola"],
},
"geo_index": {
"table_name": "table_with_geo",
"columns": ["geocol"],
"type": "geo",
},
"json_index": {
"table_name": "table_with_json",
"columns": ["somejsoncol"],
"type": "json",
},
"existing": {}
}
self.storage._drop_indexes = MagicMock()
self.storage._get_existing_indexes = lambda: ['existing']
self.storage._init_indexes()
self.storage.engine.execute.assert_has_calls([
call(
"CREATE INDEX IF NOT EXISTS \"indexname\" ON sometable USING BTREE(cola,colb)"
),
call().close(),
call(
"CREATE INDEX IF NOT EXISTS \"index2name\" ON someothertable USING BTREE(cola)"
),
call().close(),
call(
"CREATE INDEX IF NOT EXISTS \"geo_index\" ON table_with_geo USING GIST(geocol)"
),
call().close(),
call(
"CREATE INDEX IF NOT EXISTS \"json_index\" ON table_with_json USING GIN(somejsoncol)"
),
call().close(),
])
self.storage._drop_indexes.assert_called_once()
def test_create_temporary_table(self):
expected_table = f'{self.msg["header"]["catalogue"]}_{self.msg["header"]["entity"]}_tmp'
self.storage.create_temporary_table()
for entity in self.msg["contents"]:
self.storage.write_temporary_entity(entity)
# And the engine has been called to fill the temporary table
self.storage.engine.execute.assert_called()
def test_create_temporary_table_exists(self):
expected_table = f'{self.msg["header"]["catalogue"]}_{self.msg["header"]["entity"]}_tmp'
mock_table = MagicMock()
# Make sure the test table already exists
self.storage.base.metadata.tables = {expected_table: mock_table}
expected_table = self.storage.create_temporary_table()
# Assert the drop function is called
self.storage.engine.execute.assert_any_call(
f"DROP TABLE IF EXISTS {expected_table}")
for entity in self.msg["contents"]:
self.storage.write_temporary_entity(entity)
# And the engine has been called to fill the temporary table
self.storage.engine.execute.assert_called()
def test_compare_temporary_data(self):
current = f'{self.msg["header"]["catalogue"]}_{self.msg["header"]["entity"]}'
temporary = f'{self.msg["header"]["catalogue"]}_{self.msg["header"]["entity"]}_tmp'
fields = ['_tid']
query = queries.get_comparison_query('any source', current, temporary,
fields)
diff = self.storage.compare_temporary_data(temporary)
results = [result for result in diff]
self.storage.engine.execution_options.assert_called_with(
stream_results=True)
# Assert the query is performed is deleted
self.storage.engine.execution_options().execute.assert_any_call(query)
# Assert the temporary table is deleted
self.storage.engine.execute.assert_any_call(
f"DROP TABLE IF EXISTS {temporary}")
def test_bulk_insert(self):
insert_data = {'key': 'value'}
table = MagicMock()
self.storage.bulk_insert(table, insert_data)
# Assert the query is performed
self.storage.engine.execute.assert_called()
def test_delete_confirms(self):
catalogue = self.msg["header"]["catalogue"]
entity = self.msg["header"]["entity"]
events = self.storage.EVENTS_TABLE
self.storage.delete_confirms()
self.storage.engine.execute.assert_called()
args = self.storage.engine.execute.call_args[0][0]
args = ' '.join(args.split())
expect = f"DELETE FROM {events} WHERE catalogue = '{catalogue}' AND entity = '{entity}' AND action IN ('BULKCONFIRM', 'CONFIRM')"
self.assertEqual(args, expect)
def test_flush_entities(self):
self.storage.session = MagicMock()
self.storage.FORCE_FLUSH_PER = 5
self.storage.added_session_entity_cnt = 4
self.storage._flush_entities()
self.storage.session.flush.assert_not_called()
self.assertEqual(4, self.storage.added_session_entity_cnt)
self.storage.added_session_entity_cnt = 5
self.storage._flush_entities()
self.storage.session.flush.assert_called_once()
self.assertEqual(0, self.storage.added_session_entity_cnt)
def test_get_query_value(self):
self.storage.get_query_value('SELECT * FROM test')
# Assert the query is performed
self.storage.engine.execute.assert_called_with('SELECT * FROM test')
def test_combinations_plain(self):
mock_session = MagicMock()
self.storage.get_session = mock_session
result = self.storage.get_source_catalogue_entity_combinations()
mock_session.return_value.__enter__().execute.assert_called_with(
'SELECT DISTINCT source, catalogue, entity FROM events')
def test_combinations_with_args(self):
mock_session = MagicMock()
self.storage.get_session = mock_session
result = self.storage.get_source_catalogue_entity_combinations(
col="val")
mock_session.return_value.__enter__().execute.assert_called_with(
"SELECT DISTINCT source, catalogue, entity FROM events WHERE col = 'val'"
)
def test_get_tablename(self):
self.storage.gob_model = MagicMock()
result = self.storage._get_tablename()
self.assertEqual(self.storage.gob_model.get_table_name.return_value,
result)
self.storage.gob_model.get_table_name.assert_called_with(
self.storage.metadata.catalogue, self.storage.metadata.entity)
def test_analyze_table(self):
self.storage.engine = MagicMock()
self.storage._get_tablename = lambda: 'tablename'
self.storage.analyze_table()
self.storage.engine.connect.return_value.execute.assert_called_with(
'VACUUM ANALYZE tablename')
def test_add_events(self):
self.storage.session = MagicMock()
metadata = fixtures.get_metadata_fixture()
event = fixtures.get_event_fixture(metadata, 'ADD')
event['data'] = {
'_source_id': "source_id + escape '% ",
'_tid': "abcd.1 + escape '% "
}
expected = f"""
INSERT INTO
"{self.storage.EVENTS_TABLE}"
(
"timestamp",
catalogue,
entity,
"version",
"action",
"source",
source_id,
contents,
application,
tid
)
VALUES (
'{ self.storage.metadata.timestamp }',
'meetbouten',
'meetbouten',
'0.9',
'ADD',
'{ self.storage.metadata.source }',
'source_id + escape \'\'%% ',
'{{"_source_id": "source_id + escape \'\'%% ", "_tid": "abcd.1 + escape \'\'%% "}}',
'{ self.storage.metadata.application }',
'abcd.1 + escape \'\'%% '
)"""
self.storage.add_events([event])
self.storage.engine.execute.assert_called()
args = self.storage.engine.execute.call_args[0][0]
args = ' '.join(args.split())
self.assertEqual(args, ' '.join(expected.split()))