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)