示例#1
0
    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)
示例#2
0
    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)