Supports downgrade migration.

PiperOrigin-RevId: 277809167

RELEASE.md

@@ -11,6 +11,8 @@
 *   Supports Sqlite for Windows and adds scripts to build wheels for python3 in
     Windows.
 *   Provides GetContextTypes to list all Context Types.
+*   MLMD ConnectionConfig provides an option to disable an automatic upgrade.
+*   Supports downgrade of the database schema version to older versions.
 
 ## Bug Fixes and Other Changes
 

g3doc/get_started.md

@@ -324,6 +324,56 @@ version of the MLMD library (`library_version`) with the schema version
     case, please report issues for a fix or downgrade library to work with the
     database.
 
-*   If `library_version` is older than `db_version`, MLMD library returns errors
-    to prevent any data loss. In this case, the user should upgrade the library
-    version before using that database.
+*   If `library_version` is older than `db_version`, by default MLMD library
+    returns errors to prevent any data loss. In this case, the user should
+    upgrade the library version before using that database.
+
+#### Turn-off upgrade migration during connection
+
+If the upgrade migration during connection is not appropriate in the deployment
+setting, the feature can be turned-off explicitly by setting the migration
+options `disable_upgrade_migration` when creating the metadata store. MLMD will
+only check the compatibility and raise errors when the versions are
+incompatible.
+
+For example:
+
+```python
+connection_config = metadata_store_pb2.ConnectionConfig()
+connection_config.sqlite.filename_uri = '...'
+store = metadata_store.MetadataStore(connection_config,
+                                     disable_upgrade_migration = True)
+```
+
+#### Downgrade the database schema
+
+A misconfiguration in the deployment of MLMD may cause an accidental upgrade,
+e.g., when an engineer tries out a new version of the library and accidentally
+connects to the production instance of MLMD. To recover from these situations,
+MLMD provides a downgrade feature. During connection, if the migration options
+specify the `downgrade_to_schema_version`, MLMD will run a downgrade transaction
+to revert the schema version and migrate the data, then terminate the
+connection. Since the update is transactional, any failure will cause a full
+rollback of the downgrade. Once the downgrade is done, the user needs to use the
+older version of the library to connect to the database.
+
+For example:
+
+```python
+connection_config = metadata_store_pb2.ConnectionConfig()
+connection_config.sqlite.filename_uri = '...'
+metadata_store.downgrade_schema(connection_config,
+                                downgrade_to_schema_version = 0)
+```
+
+NOTE: When downgrading, MLMD prevents data loss as much as possible. However,
+newer schema versions might be inherently more expressive and hence a downgrade
+can introduce data loss.
+
+The list of `schema_version` used in MLMD releases are:
+
+ml-metadata (MLMD) | schema_version
+------------------ | --------------
+0.15.0             | 4
+0.14.0             | 4
+0.13.2             | 0

ml_metadata/metadata_store/BUILD

@@ -94,6 +94,7 @@ ml_metadata_cc_test(
         "@com_google_absl//absl/memory",
         "//ml_metadata/proto:metadata_source_proto",
         "//ml_metadata/util:metadata_source_query_config",
+        "@org_tensorflow//tensorflow/core:lib",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core:test",
     ],

ml_metadata/metadata_store/metadata_access_object.cc

@@ -1461,7 +1461,8 @@ tensorflow::Status MetadataAccessObject::GetSchemaVersion(int64* db_version) {
   return tensorflow::errors::NotFound("it looks an empty db is given.");
 }
 
-tensorflow::Status MetadataAccessObject::UpgradeMetadataSourceIfOutOfDate() {
+tensorflow::Status MetadataAccessObject::UpgradeMetadataSourceIfOutOfDate(
+    const bool disable_migration) {
   const int64 lib_version = query_config_.schema_version();
   int64 db_version = 0;
   tensorflow::Status get_schema_version_status = GetSchemaVersion(&db_version);
@@ -1471,14 +1472,27 @@ tensorflow::Status MetadataAccessObject::UpgradeMetadataSourceIfOutOfDate() {
   } else {
     TF_RETURN_IF_ERROR(get_schema_version_status);
   }
-  // we don't support downgrade a live database.
+  // we don't know how to downgrade a newer live database to the current
+  // library version as the newer schema is in a later version of the library.
   if (db_version > lib_version) {
     return tensorflow::errors::FailedPrecondition(
         "MLMD database version ", db_version,
         " is greater than library version ", lib_version,
         ". Please upgrade the library to use the given database in order to "
-        "prevent potential data loss.");
+        "prevent potential data loss. If data loss is acceptable, please"
+        " downgrade the database using a newer version of library.");
   }
+  // returns error if upgrade is explicitly disabled, as we are missing schema
+  // and cannot continue with this library version.
+  if (db_version < lib_version && disable_migration) {
+    return tensorflow::errors::FailedPrecondition(
+        "MLMD database version ", db_version, " is older than library version ",
+        lib_version,
+        ". Schema migration is disabled. Please upgrade the database then use"
+        " the library version; or switch to a older library version to use the"
+        " current database.");
+  }
+
   // migrate db_version to lib version
   const auto& migration_schemes = query_config_.migration_schemes();
   while (db_version < lib_version) {
@@ -1497,19 +1511,73 @@ tensorflow::Status MetadataAccessObject::UpgradeMetadataSourceIfOutOfDate() {
         ComposeParameterizedQuery(query_config_.update_schema_version(),
                                   {Bind(to_version)}, &update_schema_version));
     upgrade_queries.push_back(update_schema_version);
-    std::vector<RecordSet> dummy_record_sets;
     TF_RETURN_WITH_CONTEXT_IF_ERROR(
-        ExecuteMultiQuery(upgrade_queries, metadata_source_,
-                          &dummy_record_sets),
+        ExecuteMultiQuery(upgrade_queries, metadata_source_),
         "Failed to migrate existing db; the migration transaction rolls back.");
     db_version = to_version;
   }
   return tensorflow::Status::OK();
 }
 
-tensorflow::Status MetadataAccessObject::InitMetadataSourceIfNotExists() {
+tensorflow::Status MetadataAccessObject::DowngradeMetadataSource(
+    const int64 to_schema_version) {
+  const int64 lib_version = query_config_.schema_version();
+  if (to_schema_version < 0 || to_schema_version > lib_version) {
+    return tensorflow::errors::InvalidArgument(
+        "MLMD cannot be downgraded to schema_version: ", to_schema_version,
+        ". The target version should be greater or equal to 0, and the current"
+        " library version: ",
+        lib_version, " needs to be greater than the target version.");
+  }
+  int64 db_version = 0;
+  tensorflow::Status get_schema_version_status = GetSchemaVersion(&db_version);
+  // if it is an empty database, then we skip downgrade and returns.
+  if (tensorflow::errors::IsNotFound(get_schema_version_status)) {
+    return tensorflow::errors::InvalidArgument(
+        "Empty database is given. Downgrade operation is not needed.");
+  }
+  TF_RETURN_IF_ERROR(get_schema_version_status);
+  if (db_version > lib_version) {
+    return tensorflow::errors::FailedPrecondition(
+        "MLMD database version ", db_version,
+        " is greater than library version ", lib_version,
+        ". The current library does not know how to downgrade it. "
+        "Please upgrade the library to downgrade the schema.");
+  }
+  // perform downgrade
+  const auto& migration_schemes = query_config_.migration_schemes();
+  while (db_version > to_schema_version) {
+    const int64 to_version = db_version - 1;
+    if (migration_schemes.find(to_version) == migration_schemes.end()) {
+      return tensorflow::errors::Internal(
+          "Cannot find migration_schemes to version ", to_version);
+    }
+    std::vector<Query> migration_queries;
+    for (const MetadataSourceQueryConfig::TemplateQuery& downgrade_query :
+         migration_schemes.at(to_version).downgrade_queries()) {
+      migration_queries.push_back(downgrade_query.query());
+    }
+    // at version 0, v0.13.2, there is no schema version information.
+    if (to_version > 0) {
+      Query update_schema_version;
+      TF_RETURN_IF_ERROR(ComposeParameterizedQuery(
+          query_config_.update_schema_version(), {Bind(to_version)},
+          &update_schema_version));
+      migration_queries.push_back(update_schema_version);
+    }
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(
+        ExecuteMultiQuery(migration_queries, metadata_source_),
+        "Failed to migrate existing db; the migration transaction rolls back.");
+    db_version = to_version;
+  }
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status MetadataAccessObject::InitMetadataSourceIfNotExists(
+    const bool disable_upgrade_migration) {
   // check db version, and make it to align with the lib version.
-  TF_RETURN_IF_ERROR(UpgradeMetadataSourceIfOutOfDate());
+  TF_RETURN_IF_ERROR(
+      UpgradeMetadataSourceIfOutOfDate(disable_upgrade_migration));
 
   // if lib and db versions align, we check the required tables for the lib.
   const Query& check_type_table = query_config_.check_type_table().query();

ml_metadata/metadata_store/metadata_access_object.h

@@ -82,8 +82,19 @@ class MetadataAccessObject {
   // Returns OK and does nothing, if all required schema exist.
   // Returns OK and creates schema, if no schema exists yet.
   // Returns DATA_LOSS error, if any required schema is missing.
+  // Returns FAILED_PRECONDITION error, if library and db have incompatible
+  //   schema versions, and upgrade migrations are disallowed.
   // Returns detailed INTERNAL error, if create schema query execution fails.
-  tensorflow::Status InitMetadataSourceIfNotExists();
+  tensorflow::Status InitMetadataSourceIfNotExists(
+      bool disable_upgrade_migration = false);
+
+  // Downgrades the schema to `to_schema_version` in the given metadata source.
+  // Returns INVALID_ARGUMENT, if `to_schema_version` is less than 0, or newer
+  //   than the library version.
+  // Returns FAILED_PRECONDITION, if db schema version is newer than the
+  //   library version.
+  // Returns detailed INTERNAL error, if query execution fails.
+  tensorflow::Status DowngradeMetadataSource(int64 to_schema_version);
 
   // Creates a type, returns the assigned type id. A type is one of
   // {ArtifactType, ExecutionType, ContextType}. The id field of the given type
@@ -330,16 +341,19 @@ class MetadataAccessObject {
   // Upgrades the database schema version (db_v) to align with the library
   // schema version (lib_v). It retrieves db_v from the metadata source and
   // compares it with the lib_v in the given query_config, and runs migration
-  // queries if db_v < lib_v.
+  // queries if db_v < lib_v. If `disable_migration`, it only compares the
+  // db_v with lib_v and does not change the db schema.
   // Returns FAILED_PRECONDITION error, if db_v > lib_v for the case that the
   //   user use a database produced by a newer version of the library. In that
   //   case, downgrading the database may result in data loss. Often upgrading
   //   the library is required.
+  // Returns FAILED_PRECONDITION error, if db_v < lib_v and `disable_migration`
+  //   is set to true.
   // Returns DATA_LOSS error, if schema version table exists but no value found.
   // Returns DATA_LOSS error, if the database is not a 0.13.2 release database
   //   and the schema version cannot be resolved.
   // Returns detailed INTERNAL error, if query execution fails.
-  tensorflow::Status UpgradeMetadataSourceIfOutOfDate();
+  tensorflow::Status UpgradeMetadataSourceIfOutOfDate(bool disable_migration);
 
   const MetadataSourceQueryConfig query_config_;
 

ml_metadata/metadata_store/metadata_access_object_test.cc

@@ -1586,6 +1586,13 @@ TEST_P(MetadataAccessObjectTest, MigrateToCurrentLibVersion) {
     TF_EXPECT_OK(metadata_access_object_->GetSchemaVersion(&v0_13_2_version));
     EXPECT_EQ(0, v0_13_2_version);
   }
+  // expect to have error when connecting an older database version without
+  // allowing upgrade migration
+  tensorflow::Status status =
+      metadata_access_object_->InitMetadataSourceIfNotExists(
+          /*disable_upgrade_migration=*/true);
+  EXPECT_EQ(status.code(), tensorflow::error::FAILED_PRECONDITION);
+
   // then init the store and the migration queries runs.
   TF_EXPECT_OK(metadata_access_object_->InitMetadataSourceIfNotExists());
   // at the end state, schema version should becomes the library version and
@@ -1611,6 +1618,58 @@ TEST_P(MetadataAccessObjectTest, MigrateToCurrentLibVersion) {
   }
 }
 
+TEST_P(MetadataAccessObjectTest, DowngradeToV0FromCurrentLibVersion) {
+  // should not use downgrade when the database is empty.
+  EXPECT_EQ(metadata_access_object_
+                ->DowngradeMetadataSource(
+                    /*to_schema_version=*/0)
+                .code(),
+            tensorflow::error::INVALID_ARGUMENT);
+  // init the database to the current library version.
+  TF_EXPECT_OK(metadata_access_object_->InitMetadataSourceIfNotExists());
+  const int64 lib_version =
+      metadata_access_object_->query_config().schema_version();
+  int64 curr_version = 0;
+  TF_EXPECT_OK(metadata_access_object_->GetSchemaVersion(&curr_version));
+  EXPECT_EQ(curr_version, lib_version);
+
+  // downgrade one version at a time and verify the state.
+  for (int i = lib_version - 1; i >= 0; i--) {
+    // set the pre-migration states of i+1 version.
+    ASSERT_TRUE(
+        metadata_access_object_->query_config().migration_schemes().find(i) !=
+        metadata_access_object_->query_config().migration_schemes().end());
+    const MetadataSourceQueryConfig::MigrationScheme scheme =
+        metadata_access_object_->query_config().migration_schemes().at(i);
+    if (!scheme.has_downgrade_verification()) continue;
+    for (const MetadataSourceQueryConfig::TemplateQuery& setup_query :
+         scheme.downgrade_verification().previous_version_setup_queries()) {
+      RecordSet dummy_record_set;
+      TF_EXPECT_OK(metadata_access_object_->metadata_source()->ExecuteQuery(
+          setup_query.query(), &dummy_record_set));
+    }
+
+    // downgrade
+    TF_ASSERT_OK(metadata_access_object_->DowngradeMetadataSource(i));
+
+    // verify the state of the schema
+    for (const MetadataSourceQueryConfig::TemplateQuery& verification_query :
+         scheme.downgrade_verification()
+             .post_migration_verification_queries()) {
+      RecordSet record_set;
+      TF_EXPECT_OK(metadata_access_object_->metadata_source()->ExecuteQuery(
+          verification_query.query(), &record_set));
+      ASSERT_EQ(record_set.records_size(), 1);
+      bool result = false;
+      ASSERT_TRUE(absl::SimpleAtob(record_set.records(0).values(0), &result));
+      EXPECT_TRUE(result);
+    }
+    // verify the db schema version
+    TF_EXPECT_OK(metadata_access_object_->GetSchemaVersion(&curr_version));
+    EXPECT_EQ(curr_version, i);
+  }
+}
+
 }  // namespace
 }  // namespace testing
 }  // namespace ml_metadata

ml_metadata/metadata_store/metadata_store.cc

@@ -94,10 +94,13 @@ tensorflow::Status MetadataStore::InitMetadataStore() {
       });
 }
 
-tensorflow::Status MetadataStore::InitMetadataStoreIfNotExists() {
+tensorflow::Status MetadataStore::InitMetadataStoreIfNotExists(
+    const bool disable_upgrade_migration) {
   return ExecuteTransaction(
-      metadata_source_.get(), [this]() -> tensorflow::Status {
-        return metadata_access_object_->InitMetadataSourceIfNotExists();
+      metadata_source_.get(),
+      [this, &disable_upgrade_migration]() -> tensorflow::Status {
+        return metadata_access_object_->InitMetadataSourceIfNotExists(
+            disable_upgrade_migration);
       });
 }
 
@@ -413,11 +416,27 @@ tensorflow::Status MetadataStore::PutContexts(const PutContextsRequest& request,
 
 tensorflow::Status MetadataStore::Create(
     const MetadataSourceQueryConfig& query_config,
+    const MigrationOptions& migration_options,
     unique_ptr<MetadataSource> metadata_source,
     unique_ptr<MetadataStore>* result) {
   unique_ptr<MetadataAccessObject> metadata_access_object;
   TF_RETURN_IF_ERROR(MetadataAccessObject::Create(
       query_config, metadata_source.get(), &metadata_access_object));
+  // if downgrade migration is specified
+  if (migration_options.downgrade_to_schema_version() >= 0) {
+    TF_RETURN_IF_ERROR(ExecuteTransaction(
+        metadata_source.get(),
+        [&migration_options, &metadata_access_object]() -> tensorflow::Status {
+          return metadata_access_object->DowngradeMetadataSource(
+              migration_options.downgrade_to_schema_version());
+        }));
+    return tensorflow::errors::Cancelled(
+        "Downgrade migration was performed. Connection to the downgraded "
+        "database is Cancelled. Now the database is at schema version ",
+        migration_options.downgrade_to_schema_version(),
+        ". Please refer to the migration guide and use lower version of the "
+        "library to connect to the metadata store.");
+  }
   *result = absl::WrapUnique(new MetadataStore(
       std::move(metadata_source), std::move(metadata_access_object)));
   return tensorflow::Status::OK();

ml_metadata/metadata_store/metadata_store.h

@@ -30,16 +30,17 @@ namespace ml_metadata {
 // Each method is an atomic operation.
 class MetadataStore {
  public:
-  // Factory method, if the return value is ok, 'result' is populated with an
-  // object that can be used to access metadata with the given config and
-  // metadata_source.
+  // Factory method that creates a MetadataStore in result. The result is owned
+  // by the caller, and metadata_source is owned by result.
+  // If the return value is ok, 'result' is populated with an object that can be
+  // used to access metadata with the given config and metadata_source.
   // Returns INVALID_ARGUMENT error, if query_config is not valid.
+  // Returns INVALID_ARGUMENT error, if migration options are invalid.
+  // Returns CANCELLED error, if downgrade migration is performed.
   // Returns detailed INTERNAL error, if the MetadataSource cannot be connected.
-  // Creates a MetadataStore in result.
-  // result is owned by the caller.
-  // metadata_source is owned by result.
   static tensorflow::Status Create(
       const MetadataSourceQueryConfig& query_config,
+      const MigrationOptions& migration_options,
       std::unique_ptr<MetadataSource> metadata_source,
       std::unique_ptr<MetadataStore>* result);
 
@@ -52,8 +53,11 @@ class MetadataStore {
   // Returns OK and does nothing, if all required schema exist.
   // Returns OK and creates schema, if no schema exists yet.
   // Returns DATA_LOSS error, if any required schema is missing.
+  // Returns FAILED_PRECONDITION error, if library and db have incompatible
+  //   schema versions, and upgrade migrations are disallowed.
   // Returns detailed INTERNAL error, if create schema query execution fails.
-  tensorflow::Status InitMetadataStoreIfNotExists();
+  tensorflow::Status InitMetadataStoreIfNotExists(
+      bool disable_upgrade_migration = false);
 
   // Inserts or updates an artifact type.
   //