absl::StatusOr SchemaUpdaterImpl::BuildForeignKeyConstraint()

in backend/schema/updater/schema_updater.cc [2418:2512]


absl::StatusOr<const ForeignKey*> SchemaUpdaterImpl::BuildForeignKeyConstraint(
    const ddl::ForeignKey& ddl_foreign_key, const Table* referencing_table) {
  ForeignKey::Builder builder;
  std::optional<uint32_t> oid = pg_oid_assigner_->GetNextPostgresqlOid();
  builder.set_postgresql_oid(oid);
  if (oid.has_value()) {
    ZETASQL_VLOG(2) << "Assigned oid " << oid.value() << " for FOREIGN KEY constraint "
            << ddl_foreign_key.constraint_name() << " on table "
            << referencing_table->Name();
  }

  ZETASQL_RETURN_IF_ERROR(
      AlterNode<Table>(referencing_table, [&](Table::Editor* editor) {
        editor->add_foreign_key(builder.get());
        return absl::OkStatus();
      }));
  builder.set_referencing_table(referencing_table);

  const Table* referenced_table = latest_schema_->FindTableCaseSensitive(
      ddl_foreign_key.referenced_table_name());
  if (referenced_table == nullptr) {
    if (ddl_foreign_key.referenced_table_name() != referencing_table->Name()) {
      return error::TableNotFound(ddl_foreign_key.referenced_table_name());
    }
    // Self-referencing foreign key.
    referenced_table = referencing_table;
  }
  ZETASQL_RETURN_IF_ERROR(
      AlterNode<Table>(referenced_table, [&](Table::Editor* editor) {
        editor->add_referencing_foreign_key(builder.get());
        return absl::OkStatus();
      }));
  builder.set_referenced_table(referenced_table);

  std::string foreign_key_name;
  if (ddl_foreign_key.has_constraint_name()) {
    foreign_key_name = ddl_foreign_key.constraint_name();
    ZETASQL_RETURN_IF_ERROR(global_names_.AddName("Foreign Key", foreign_key_name));
    builder.set_constraint_name(foreign_key_name);
  } else {
    ZETASQL_ASSIGN_OR_RETURN(foreign_key_name,
                     global_names_.GenerateForeignKeyName(
                         referencing_table->Name(), referenced_table->Name()));
    builder.set_generated_name(foreign_key_name);
  }

  auto add_columns =
      [&](const Table* table,
          const google::protobuf::RepeatedPtrField<std::string>& column_names,
          std::function<void(const Column*)> add_column) {
        for (const std::string& column_name : column_names) {
          const Column* column = table->FindColumnCaseSensitive(column_name);
          if (column == nullptr) {
            return error::ForeignKeyColumnNotFound(column_name, table->Name(),
                                                   foreign_key_name);
          }
          add_column(column);
        }
        return absl::OkStatus();
      };
  ZETASQL_RETURN_IF_ERROR(add_columns(referencing_table,
                              ddl_foreign_key.constrained_column_name(),
                              [&builder](const Column* column) {
                                builder.add_referencing_column(column);
                              }));
  ZETASQL_RETURN_IF_ERROR(add_columns(referenced_table,
                              ddl_foreign_key.referenced_column_name(),
                              [&builder](const Column* column) {
                                builder.add_referenced_column(column);
                              }));
  if (ddl_foreign_key.has_on_delete()) {
    if (ddl_foreign_key.on_delete() != ddl::ForeignKey::ACTION_UNSPECIFIED &&
        ddl_foreign_key.on_delete() != ddl::ForeignKey::NO_ACTION &&
        !EmulatorFeatureFlags::instance()
             .flags()
             .enable_fk_delete_cascade_action) {
      return error::ForeignKeyOnDeleteActionUnsupported(
          ForeignKey::ActionName(GetForeignKeyOnDeleteAction(ddl_foreign_key)));
    }
    builder.set_delete_action(GetForeignKeyOnDeleteAction(ddl_foreign_key));
  }

  if (!ddl_foreign_key.enforced()) {
    if (!EmulatorFeatureFlags::instance()
             .flags()
             .enable_fk_enforcement_option) {
      return error::ForeignKeyEnforcementUnsupported();
    }
    builder.set_enforced(ddl_foreign_key.enforced());
  }

  const ForeignKey* foreign_key = builder.get();
  ZETASQL_RETURN_IF_ERROR(AddNode(builder.build()));
  return foreign_key;
}