backend/schema/updater/schema_validation_context.h

//
// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#ifndef THIRD_PARTY_CLOUD_SPANNER_EMULATOR_BACKEND_SCHEMA_CATALOG_SCHEMA_VALIDATION_CONTEXT_H_
#define THIRD_PARTY_CLOUD_SPANNER_EMULATOR_BACKEND_SCHEMA_CATALOG_SCHEMA_VALIDATION_CONTEXT_H_

#include <memory>
#include <vector>

#include "zetasql/public/types/type_factory.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/time/time.h"
#include "backend/query/analyzer_options.h"
#include "backend/schema/catalog/proto_bundle.h"
#include "backend/schema/graph/schema_graph.h"
#include "backend/schema/graph/schema_node.h"
#include "backend/storage/storage.h"
#include "absl/status/status.h"
#include "zetasql/base/status_macros.h"

namespace google {
namespace spanner {
namespace emulator {
namespace backend {

class Schema;
class GlobalSchemaNames;

// A class used to collect and execute verification/backfill actions resulting
// from a schema change. A `SchemaChangeAction` object can be constructed and
// added to this context during processing/validation of schema changes that
// requires verifying that the data in the database conforms to an
// existing/newly added constraint/invariant or backfilling values in the
// database.
// Note: Currently SchemaChangeActions that modify the state of a database
// cannot be rolled back. Consequently, action execution should be atomic, i.e.
// the action should either apply all the effects entailed by the schema change
// or leave the datatabse unchanged.
class SchemaValidationContext {
 public:
  // The actual function doing the verification to which `context_` is passed.
  using SchemaChangeAction =
      std::function<absl::Status(const SchemaValidationContext*)>;

  // A callback used to return an unowned instance of a Schema.
  using SchemaConstructorCb = std::function<const Schema*(const SchemaGraph*)>;

  // Test-only constructor.
  // TODO : Split out into a separate Schema update validation
  // context that is used by SchemaUpdater and can therefore be mocked
  // separately.
  SchemaValidationContext() = default;

  SchemaValidationContext(Storage* storage, GlobalSchemaNames* global_names,
                          zetasql::TypeFactory* type_factory,
                          absl::Time pending_commit_timestamp,
                          DatabaseDialect dialect)
      : storage_(storage),
        global_names_(global_names),
        type_factory_(type_factory),
        pending_commit_timestamp_(pending_commit_timestamp),
        dialect_(dialect) {}

  ~SchemaValidationContext() = default;

  // TODO : Split out into a separate StatementValidationContext.
  // Interface accessed by validators to enqueue schema
  // change actions.
  // -----------------------------------------------------

  // Adds a SchemaChangeAction to this validation context.
  void AddAction(SchemaChangeAction action_fn) {
    actions_.emplace_back(std::move(action_fn));
  }

  // Interface used by a SchemaChangeAction to access the
  // database
  // --------------------------------------------------

  // Provides SchemaVerifiers running under this validation context
  // access to the database storage to perform reads.
  Storage* storage() const { return storage_; }

  // Access to the current schema's set of global names.
  GlobalSchemaNames* global_names() const { return global_names_; }

  // Access to the current database's type factory.
  zetasql::TypeFactory* type_factory() const { return type_factory_; }

  // Returns the pending commit timestamp.
  absl::Time pending_commit_timestamp() const {
    return pending_commit_timestamp_;
  }

  // Returns the Schema snapshot for the old/unmodified schema.
  const Schema* old_schema() const { return old_schema_snapshot_; }

  // Returns the Schema snapshot for the updated validated schema. This is only
  // valid to call during the verification and backfill phase. Attempts to call
  // this during the validation phase will result in undefined behavior. Callers
  // wishing to access a temporary, unvalidated snapshot of the new schema
  // during the validation phase should call `tmp_new_schema()` instead.
  const Schema* validated_new_schema() const { return new_schema_snapshot_; }

  // Returns the temporary schema snapshot of the new schema during the
  // validation phase and is invalid to call during the verification and
  // backfill phases. Callers should not hold on to any references to the
  // returned schema or its nodes.
  const Schema* tmp_new_schema() const { return tmp_new_schema_; }

  bool is_postgresql_dialect() const {
    return dialect_ == DatabaseDialect::POSTGRESQL;
  }

  // Interface accessed by SchemaUpdater to execute queued
  // actions.
  // -----------------------------------------------------

  void SetOldSchemaSnapshot(const Schema* old_schema) {
    old_schema_snapshot_ = old_schema;
  }

  void SetValidatedNewSchemaSnapshot(const Schema* new_schema) {
    new_schema_snapshot_ = new_schema;
  }

  void SetTempNewSchemaSnapshotConstructor(
      SchemaConstructorCb schema_constructor) {
    tmp_new_schema_cb_ = std::move(schema_constructor);
  }

  // Runs all SchemaVerifiers added to this validation context.
  absl::Status RunSchemaChangeActions() const {
    for (auto& action : actions_) {
      ZETASQL_RETURN_IF_ERROR(action(this));
    }
    return absl::OkStatus();
  }

  // Returns the number of pending schema change actions.
  int num_actions() const { return actions_.size(); }

  // Returns true if 'node' is a node that was modified using a DDL
  // statement/operation as a part of the schema change associated
  // with this SchemaValidationContext.
  bool IsModifiedNode(const SchemaNode* node) const {
    return edited_nodes_->contains(node);
  }

  // Returns the pointer to the first node added with a given name
  // (case-insensitive) using the node's SchemaNameInfo. Returns nullptr if not
  // found. Also returns nullptr if nodes with that name were found but were the
  // wrong type.
  template <typename T>
  const T* FindAddedNode(absl::string_view name) const {
    for (const auto& node : *added_nodes_) {
      auto info = node->GetSchemaNameInfo().value_or(SchemaNameInfo{});
      if (info.name == name) {
        const T* candidate = node.get()->As<T>();
        if (info.global || candidate != nullptr) {
          return candidate;
        }
      }
    }
    return nullptr;
  }

  // This method is used to set the recently parsed Proto Bundle to schema
  // validation context. This will be used during the validation of columns.
  void set_proto_bundle(std::shared_ptr<const ProtoBundle> proto_bundle) {
    proto_bundle_ = proto_bundle;
  }

  const ProtoBundle* proto_bundle() const { return proto_bundle_.get(); }

 private:
  friend class SchemaGraphEditor;

  // Saves a pointer to the 'edited_nodes' maintained by the 'SchemaGraphEditor'
  // to allow nodes to check through the SchemaValidationContext if they were
  // modified by user action or just cloned during the processing of a schema
  // update.
  void set_edited_nodes(
      const absl::flat_hash_set<const SchemaNode*>* edited_nodes) {
    edited_nodes_ = edited_nodes;
  }

  const absl::flat_hash_set<const SchemaNode*>* edited_nodes_ = nullptr;

  // Saves a pointer to the 'added_nodes' maintained by the 'SchemaGraphEditor'
  // to allow nodes to check through the SchemaValidationContext if they were
  // added by user action.
  void set_added_nodes(
      const std::vector<std::unique_ptr<const SchemaNode>>* added_nodes) {
    added_nodes_ = added_nodes;
  }

  // Called by SchemaGraphEditor to expose a temporary `schema_graph` via
  // the `Schema` interface to Validate/ValidateUpdate() functions during the
  // validation phase.
  void MakeNewTempSchemaSnapshot(const SchemaGraph* schema_graph) {
    if (tmp_new_schema_cb_ == nullptr) return;
    tmp_new_schema_ = tmp_new_schema_cb_(schema_graph);
  }

  void ClearNewTempSchemaSnapshot() { tmp_new_schema_ = nullptr; }

  const std::vector<std::unique_ptr<const SchemaNode>>* added_nodes_ = nullptr;

  // Used to read data from the database for verifiers. Not owned.
  Storage* storage_;

  // Global names for schema objects.
  GlobalSchemaNames* global_names_;

  // Type factory used for all ZetaSQL operations on the database.
  zetasql::TypeFactory* type_factory_;

  // Planned commit time for the schema change.
  absl::Time pending_commit_timestamp_;

  // The database dialect determines how the OIDs are validated.
  DatabaseDialect dialect_;

  // The list of pending schema change actions (verifications/backfills) to run.
  std::vector<SchemaChangeAction> actions_;

  // The old schema.
  const Schema* old_schema_snapshot_ = nullptr;

  // The new schema.
  const Schema* new_schema_snapshot_ = nullptr;

  // Callback to construct a temporary instance of the new Schema.
  SchemaConstructorCb tmp_new_schema_cb_ = nullptr;

  // Holds an instance of the temporary new schema during the validation phase.
  // This instance is not owned by SchemaValidationContext but is guaranteed
  // to be alive during the validation phase.
  const Schema* tmp_new_schema_ = nullptr;

  // Proto bundle for schema validation.
  // Use this only during schema validation when the schema isn't created. Once
  // the validation is completed, a schema object is created and has the proto
  // bundle which can be accessed by schema->proto_bundle()
  std::shared_ptr<const ProtoBundle> proto_bundle_;
};

}  // namespace backend
}  // namespace emulator
}  // namespace spanner
}  // namespace google

#endif  // THIRD_PARTY_CLOUD_SPANNER_EMULATOR_BACKEND_SCHEMA_CATALOG_SCHEMA_VALIDATION_CONTEXT_H_