def test_enable_metadata_store_upgrade_migration(self):
# create a metadata store and downgrade to version 0
db_file = os.path.join(absltest.get_default_test_tmpdir(), "test.db")
connection_config = metadata_store_pb2.ConnectionConfig()
connection_config.sqlite.filename_uri = db_file
metadata_store.MetadataStore(connection_config)
metadata_store.downgrade_schema(connection_config, 0)
upgrade_conn_config = metadata_store_pb2.ConnectionConfig()
upgrade_conn_config.sqlite.filename_uri = db_file
with self.assertRaisesRegex(RuntimeError,
"chema migration is disabled"):
# if disabled then the store cannot be used.
metadata_store.MetadataStore(upgrade_conn_config)
# if enable, then the store can be created
metadata_store.MetadataStore(upgrade_conn_config,
enable_upgrade_migration=True)
os.remove(db_file)
def test_downgrade_metadata_store(self):
# create a metadata store and init to the current library version
connection_config = metadata_store_pb2.ConnectionConfig()
db_file = os.path.join(absltest.get_default_test_tmpdir(), "test.db")
connection_config.sqlite.filename_uri = db_file
metadata_store.MetadataStore(connection_config)
# wrong downgrade_to_schema_version
with self.assertRaisesRegex(
errors.InvalidArgumentError,
"downgrade_to_schema_version not specified"):
metadata_store.downgrade_schema(connection_config, -1)
# invalid argument for the downgrade_to_schema_version
with self.assertRaisesRegex(
errors.InvalidArgumentError,
"MLMD cannot be downgraded to schema_version"):
downgrade_to_version = 999999
metadata_store.downgrade_schema(connection_config,
downgrade_to_version)
# downgrade the metadata store to v0.13.2 where schema version is 0
metadata_store.downgrade_schema(connection_config,
downgrade_to_schema_version=0)
os.remove(db_file)
