backend/schema/updater/schema_updater_tests/identity_column.cc

//
// 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.
//

#include <memory>
#include <optional>

#include "google/spanner/admin/database/v1/common.pb.h"
#include "zetasql/public/options.pb.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "zetasql/base/testing/status_matchers.h"
#include "tests/common/proto_matchers.h"
#include "absl/status/status.h"
#include "backend/schema/catalog/column.h"
#include "backend/schema/catalog/schema.h"
#include "backend/schema/catalog/sequence.h"
#include "backend/schema/updater/schema_updater_tests/base.h"

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

using database_api::DatabaseDialect::POSTGRESQL;
using ::testing::HasSubstr;
using ::zetasql_base::testing::StatusIs;
using IdentityColumnTest = SchemaUpdaterTest;

INSTANTIATE_TEST_SUITE_P(
    SchemaUpdaterPerDialectTests, IdentityColumnTest,
    testing::Values(database_api::DatabaseDialect::GOOGLE_STANDARD_SQL,
                    database_api::DatabaseDialect::POSTGRESQL),
    [](const testing::TestParamInfo<IdentityColumnTest::ParamType>& info) {
      return database_api::DatabaseDialect_Name(info.param);
    });

TEST_P(IdentityColumnTest, Basic) {
  std::unique_ptr<const Schema> schema;
  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
    CREATE TABLE T (
      id INT64 GENERATED BY DEFAULT AS IDENTITY (
        BIT_REVERSED_POSITIVE START COUNTER WITH 10 SKIP RANGE 100, 1000),
      non_key_col INT64 GENERATED BY DEFAULT AS IDENTITY (
        BIT_REVERSED_POSITIVE START COUNTER WITH 20 SKIP RANGE 200, 2000),
      value INT64,
    ) PRIMARY KEY(id)
      )"}));

  const Table* table = schema->FindTable("T");
  ASSERT_NE(table, nullptr);
  const Column* col = table->FindColumn("id");
  ASSERT_NE(col, nullptr);
  EXPECT_TRUE(col->has_default_value());
  EXPECT_TRUE(col->is_identity_column());
  ASSERT_TRUE(col->expression().has_value());
  if (GetParam() == POSTGRESQL) {
    EXPECT_EQ(col->expression().value(),
              "(GET_NEXT_SEQUENCE_VALUE(\"_identity_seq_T__id\"))");
  } else {
    EXPECT_EQ(col->expression().value(),
              "(GET_NEXT_SEQUENCE_VALUE(SEQUENCE _identity_seq_T__id))");
  }

  ASSERT_EQ(col->sequences_used().size(), 1);
  const Sequence* seq =
      static_cast<const Sequence*>(col->sequences_used().at(0));
  ASSERT_NE(seq, nullptr);
  EXPECT_EQ(seq->Name(), "_identity_seq_T__id");
  EXPECT_EQ(seq->sequence_kind_name(), "BIT_REVERSED_POSITIVE");
  EXPECT_EQ(seq->start_with_counter(), 10);
  EXPECT_EQ(seq->skip_range_min(), 100);
  EXPECT_EQ(seq->skip_range_max(), 1000);

  col = table->FindColumn("non_key_col");
  ASSERT_NE(col, nullptr);
  EXPECT_TRUE(col->has_default_value());
  EXPECT_TRUE(col->is_identity_column());
  ASSERT_TRUE(col->expression().has_value());
  if (GetParam() == POSTGRESQL) {
    EXPECT_EQ(col->expression().value(),
              "(GET_NEXT_SEQUENCE_VALUE(\"_identity_seq_T__non_key_col\"))");
  } else {
    EXPECT_EQ(
        col->expression().value(),
        "(GET_NEXT_SEQUENCE_VALUE(SEQUENCE _identity_seq_T__non_key_col))");
  }

  ASSERT_EQ(col->sequences_used().size(), 1);
  seq = static_cast<const Sequence*>(col->sequences_used().at(0));
  ASSERT_NE(seq, nullptr);
  EXPECT_EQ(seq->Name(), "_identity_seq_T__non_key_col");
  EXPECT_EQ(seq->sequence_kind_name(), "BIT_REVERSED_POSITIVE");
  EXPECT_EQ(seq->start_with_counter(), 20);
  EXPECT_EQ(seq->skip_range_min(), 200);
  EXPECT_EQ(seq->skip_range_max(), 2000);
}

TEST_P(IdentityColumnTest, NamedSchema) {
  std::unique_ptr<const Schema> schema;
  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
    CREATE SCHEMA sch
  )",
                                             R"(
    CREATE TABLE sch.T (
      id INT64 GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
      non_key_col INT64 GENERATED BY DEFAULT AS IDENTITY (
        BIT_REVERSED_POSITIVE),
      value INT64,
    ) PRIMARY KEY(id)
      )"}));
  const Table* table = schema->FindTable("sch.T");
  ASSERT_NE(table, nullptr);
  const Column* col = table->FindColumn("id");
  ASSERT_NE(col, nullptr);
  EXPECT_TRUE(col->is_identity_column());
  col = table->FindColumn("non_key_col");
  ASSERT_NE(col, nullptr);
  EXPECT_TRUE(col->is_identity_column());
}

TEST_P(IdentityColumnTest, DropDefault) {
  std::unique_ptr<const Schema> schema;
  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
    CREATE TABLE T (
      id INT64 GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
      value INT64,
    ) PRIMARY KEY(id)
      )",
                                             R"(
    ALTER TABLE T ALTER COLUMN id DROP DEFAULT
                                              )"}));
  const Table* table = schema->FindTable("T");
  ASSERT_NE(table, nullptr);
  const Column* col = table->FindColumn("id");
  ASSERT_NE(col, nullptr);
  EXPECT_TRUE(col->has_default_value());
  EXPECT_TRUE(col->is_identity_column());
  ASSERT_EQ(col->sequences_used().size(), 1);
  const Sequence* seq =
      static_cast<const Sequence*>(col->sequences_used().at(0));
  ASSERT_NE(seq, nullptr);
  EXPECT_EQ(seq->Name(), "_identity_seq_T__id");
}

TEST_P(IdentityColumnTest, DropTable) {
  std::unique_ptr<const Schema> schema;
  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({
                                   R"(
    CREATE TABLE T (
      id INT64 GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
      non_key_col INT64 GENERATED BY DEFAULT AS IDENTITY (
        BIT_REVERSED_POSITIVE),
      value INT64,
    ) PRIMARY KEY(id)
      )",
                                   R"(
    DROP TABLE T
                                              )"}));
  const Table* table = schema->FindTable("T");
  ASSERT_EQ(table, nullptr);
  const Sequence* seq = schema->FindSequence("_identity_seq_T__non_key_col");
  EXPECT_EQ(seq, nullptr);

  // Test to recreate the same table again.
  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, UpdateSchema(schema.get(), {
                                                              R"(
    CREATE TABLE T (
      id INT64 GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
      non_key_col INT64 GENERATED BY DEFAULT AS IDENTITY (
        BIT_REVERSED_POSITIVE),
      value INT64,
    ) PRIMARY KEY(id)
      )"}));
  table = schema->FindTable("T");
  ASSERT_NE(table, nullptr);
  seq = schema->FindSequence("_identity_seq_T__non_key_col");
  EXPECT_NE(seq, nullptr);
}

TEST_P(IdentityColumnTest, DropIdentityColumn) {
  std::unique_ptr<const Schema> schema;
  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
    CREATE TABLE T (
      id INT64 GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
      non_key_col INT64 GENERATED BY DEFAULT AS IDENTITY (
        BIT_REVERSED_POSITIVE),
      value INT64,
    ) PRIMARY KEY(id)
      )",
                                             R"(
    ALTER TABLE T DROP COLUMN non_key_col
                                              )"}));
  const Table* table = schema->FindTable("T");
  ASSERT_NE(table, nullptr);
  const Column* col = table->FindColumn("non_key_col");
  ASSERT_EQ(col, nullptr);
  const Sequence* seq = schema->FindSequence("_identity_seq_T__non_key_col");
  EXPECT_EQ(seq, nullptr);
}

TEST_P(IdentityColumnTest, AddColumn) {
  std::unique_ptr<const Schema> schema;
  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
      CREATE TABLE T (
        id INT64 GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
        value INT64,
      ) PRIMARY KEY(id)
        )",
                                             R"(
      ALTER TABLE T ADD COLUMN non_key_col INT64
            GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE)
                                               )"}));
  const Table* table = schema->FindTable("T");
  ASSERT_NE(table, nullptr);
  const Column* col = table->FindColumn("non_key_col");
  ASSERT_NE(col, nullptr);
  EXPECT_TRUE(col->has_default_value());
  EXPECT_TRUE(col->is_identity_column());
  ASSERT_EQ(col->sequences_used().size(), 1);
  const Sequence* seq =
      static_cast<const Sequence*>(col->sequences_used().at(0));
  ASSERT_NE(seq, nullptr);
  EXPECT_EQ(seq->Name(), "_identity_seq_T__non_key_col");
}

TEST_P(IdentityColumnTest, AlterEntireColumn) {
  std::unique_ptr<const Schema> schema;
  if (GetParam() == POSTGRESQL) {
    // ALTER COLUMN dosn't support altering an entire column to identity column.
    // See: ALTER COLUMN in
    // https://cloud.google.com/spanner/docs/reference/postgresql/data-definition-language#alter_table
    GTEST_SKIP();
  } else {
    ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
      CREATE TABLE T (
        id INT64 GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
        value INT64,
      ) PRIMARY KEY(id)
        )",
                                               R"(
      ALTER TABLE T ALTER COLUMN id INT64
            GENERATED BY DEFAULT AS IDENTITY (
            BIT_REVERSED_POSITIVE START COUNTER WITH 1000)
                                               )"}));
  }
  const Sequence* seq = schema->FindSequence("_identity_seq_T__id");
  ASSERT_NE(seq, nullptr);
  EXPECT_EQ(seq->start_with_counter(), 1000);
}

TEST_P(IdentityColumnTest, AlterRestartCounter) {
  std::unique_ptr<const Schema> schema;
  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
      CREATE TABLE T (
        id INT64 GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
        value INT64,
      ) PRIMARY KEY(id)
        )",
                                             R"(
      ALTER TABLE T ALTER COLUMN id ALTER IDENTITY RESTART COUNTER WITH 1000
                                               )"}));
  const Sequence* seq = schema->FindSequence("_identity_seq_T__id");
  ASSERT_NE(seq, nullptr);
  EXPECT_EQ(seq->start_with_counter(), 1000);
}

TEST_P(IdentityColumnTest, AlterSkipRange) {
  std::unique_ptr<const Schema> schema;
  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
      CREATE TABLE T (
        id INT64 GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
        value INT64,
      ) PRIMARY KEY(id)
        )",
                                             R"(
      ALTER TABLE T ALTER COLUMN id ALTER IDENTITY SET SKIP RANGE 2000, 3000
                                               )"}));
  const Sequence* seq = schema->FindSequence("_identity_seq_T__id");
  ASSERT_NE(seq, nullptr);
  EXPECT_EQ(seq->skip_range_min(), 2000);
  EXPECT_EQ(seq->skip_range_max(), 3000);

  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, UpdateSchema(schema.get(), {R"(
      ALTER TABLE T ALTER COLUMN id ALTER IDENTITY SET NO SKIP RANGE
                                               )"}));
  seq = schema->FindSequence("_identity_seq_T__id");
  ASSERT_NE(seq, nullptr);
  EXPECT_EQ(seq->skip_range_min(), std::nullopt);
  EXPECT_EQ(seq->skip_range_max(), std::nullopt);
}

TEST_P(IdentityColumnTest, SetDefaultSequenceKindFromNullToValue) {
  std::unique_ptr<const Schema> schema;
  if (GetParam() == POSTGRESQL) {
    ZETASQL_ASSERT_OK_AND_ASSIGN(schema,
                         CreateSchema({R"(
      ALTER DATABASE db SET spanner.default_sequence_kind = DEFAULT
                                      )",
                                       R"(
      ALTER DATABASE db
      SET spanner.default_sequence_kind = 'bit_reversed_positive'
                                      )"},
                                      /*proto_descriptor_bytes=*/"",
                                      /*dialect=*/POSTGRESQL,
                                      /*use_gsql_to_pg_translation=*/false));
  } else {
    ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
      ALTER DATABASE db SET OPTIONS (default_sequence_kind = NULL)
        )",
                                               R"(
      ALTER DATABASE db SET OPTIONS (
        default_sequence_kind = 'bit_reversed_positive')
                                               )"}));
  }
}

TEST_P(IdentityColumnTest, SetDefaultSequenceKindFromValueToNull) {
  std::unique_ptr<const Schema> schema;
  if (GetParam() == POSTGRESQL) {
    EXPECT_THAT(CreateSchema({R"(
      ALTER DATABASE db
      SET spanner.default_sequence_kind = 'bit_reversed_positive'
                              )",
                              R"(
      ALTER DATABASE db SET spanner.default_sequence_kind = DEFAULT
                              )"},
                             /*proto_descriptor_bytes=*/"",
                             /*dialect=*/POSTGRESQL,
                             /*use_gsql_to_pg_translation=*/false),
                StatusIs(absl::StatusCode::kFailedPrecondition,
                         HasSubstr("The default_sequence_kind option may not "
                                   "be unset once set.")));
  } else {
    EXPECT_THAT(CreateSchema({R"(
      ALTER DATABASE db SET OPTIONS (
        default_sequence_kind = 'bit_reversed_positive')
        )",
                              R"(
      ALTER DATABASE db SET OPTIONS (default_sequence_kind = NULL)
                                               )"}),
                StatusIs(absl::StatusCode::kFailedPrecondition,
                         HasSubstr("The default_sequence_kind option may not "
                                   "be unset once set.")));
  }
}

TEST_P(IdentityColumnTest, UnsupportedDefaultSequenceKindValue) {
  std::unique_ptr<const Schema> schema;
  if (GetParam() == POSTGRESQL) {
    EXPECT_THAT(
        CreateSchema({R"(
      ALTER DATABASE db
      SET spanner.default_sequence_kind = 1
                     )"},
                     /*proto_descriptor_bytes=*/"",
                     /*dialect=*/POSTGRESQL,
                     /*use_gsql_to_pg_translation=*/false),
        StatusIs(absl::StatusCode::kFailedPrecondition,
                 HasSubstr("Incorrect datatype specified for option.")));
  } else {
    EXPECT_THAT(
        CreateSchema({R"(
      ALTER DATABASE db SET OPTIONS (default_sequence_kind = 1)
                                                 )"}),
        StatusIs(
            absl::StatusCode::kInvalidArgument,
            HasSubstr("Unexpected value for option: default_sequence_kind. "
                      "Supported option values are strings and NULL.")));
  }
}

TEST_P(IdentityColumnTest, UseDefaultSequenceKindOptionFailed) {
  std::unique_ptr<const Schema> schema;
  EXPECT_THAT(
      CreateSchema({R"(
      CREATE TABLE T (
        id INT64 GENERATED BY DEFAULT AS IDENTITY,
        value INT64,
      ) PRIMARY KEY(id)
                                               )"}),
      StatusIs(
          absl::StatusCode::kInvalidArgument,
          HasSubstr("The sequence kind of an identity column id is not "
                    "specified. Please specify the sequence kind explicitly "
                    "or set the database option `default_sequence_kind`.")));
}

TEST_P(IdentityColumnTest, UseDefaultSequenceKindOptionSuccess) {
  std::unique_ptr<const Schema> schema;
  if (GetParam() == POSTGRESQL) {
    ZETASQL_ASSERT_OK_AND_ASSIGN(schema,
                         CreateSchema({R"(
      ALTER DATABASE db
      SET spanner.default_sequence_kind = 'bit_reversed_positive'
                                      )"},
                                      /*proto_descriptor_bytes=*/"",
                                      /*dialect=*/POSTGRESQL,
                                      /*use_gsql_to_pg_translation=*/false));
  } else {
    ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
      ALTER DATABASE db SET OPTIONS (
        default_sequence_kind = 'bit_reversed_positive')
        )"}));
  }
  ZETASQL_ASSERT_OK_AND_ASSIGN(schema, UpdateSchema(schema.get(), {R"(
      CREATE TABLE T (
        id INT64 GENERATED BY DEFAULT AS IDENTITY,
        value INT64,
      ) PRIMARY KEY(id)
                                               )"}));
  const Sequence* seq = schema->FindSequence("_identity_seq_T__id");
  ASSERT_NE(seq, nullptr);
  EXPECT_EQ(seq->sequence_kind_name(), "BIT_REVERSED_POSITIVE");
}

TEST_P(IdentityColumnTest, AlterNonIdentityColumn) {
  std::unique_ptr<const Schema> schema;
  EXPECT_THAT(
      CreateSchema({R"(
      CREATE TABLE T (
        id INT64 GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
        value INT64,
      ) PRIMARY KEY(id)
                                               )",
                    R"(
      ALTER TABLE T ALTER COLUMN value ALTER IDENTITY SET NO SKIP RANGE
                                              )"}),
      StatusIs(
          absl::StatusCode::kInvalidArgument,
          HasSubstr("Column is not an identity column in table T: value")));
}

TEST_P(IdentityColumnTest, SerialAutoIncrement) {
  std::unique_ptr<const Schema> schema;
  if (GetParam() == POSTGRESQL) {
    ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema(
                                     {
                                         R"(
    ALTER DATABASE db
      SET spanner.default_sequence_kind = 'bit_reversed_positive'
                                      )",
                                         R"(
    CREATE TABLE T (
      id serial,
      value bigint,
      PRIMARY KEY(id)
    )
                                      )"},
                                     /*proto_descriptor_bytes=*/"",
                                     /*dialect=*/POSTGRESQL,
                                     /*use_gsql_to_pg_translation=*/false));
  } else {
    ZETASQL_ASSERT_OK_AND_ASSIGN(schema, CreateSchema({R"(
    ALTER DATABASE db SET OPTIONS (
      default_sequence_kind = 'bit_reversed_positive')
    )",
                                               R"(
    CREATE TABLE T (
      id INT64 AUTO_INCREMENT,
      value INT64,
    ) PRIMARY KEY(id)
        )"}));
  }

  const Table* table = schema->FindTable("T");
  ASSERT_NE(table, nullptr);
  const Column* col = table->FindColumn("id");
  ASSERT_NE(col, nullptr);
  EXPECT_TRUE(col->is_identity_column());
  if (GetParam() == POSTGRESQL) {
    EXPECT_EQ(col->is_nullable(), false);
  } else {
    EXPECT_EQ(col->is_nullable(), true);
  }
}

TEST_P(IdentityColumnTest, SerialAutoIncrementWithoutSettingDatabaseOption) {
  std::unique_ptr<const Schema> schema;
  if (GetParam() == POSTGRESQL) {
    EXPECT_THAT(
        CreateSchema({R"(
    CREATE TABLE T (
      id serial,
      value bigint,
      PRIMARY KEY(id)
    )
                                      )"},
                     /*proto_descriptor_bytes=*/"",
                     /*dialect=*/POSTGRESQL,
                     /*use_gsql_to_pg_translation=*/false),
        StatusIs(absl::StatusCode::kInvalidArgument,
                 HasSubstr("The sequence kind of an identity column id is not "
                           "specified.")));
  } else {
    EXPECT_THAT(CreateSchema({R"(
    CREATE TABLE T (
      id INT64 AUTO_INCREMENT,
      value INT64,
    ) PRIMARY KEY(id)
        )"}),
                StatusIs(absl::StatusCode::kInvalidArgument,
                         HasSubstr("The sequence kind of an identity column id "
                                   "is not specified.")));
  }
}

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