Skip to content

Commit

Permalink
Merge pull request #200 from GoogleCloudPlatform/5789B3FF9E750C15D6B7…
Browse files Browse the repository at this point in the history
…4524C048BD98

Project import generated by Copybara.
  • Loading branch information
olavloite authored Dec 13, 2024
2 parents d3e265a + 7268487 commit 7ad3d59
Show file tree
Hide file tree
Showing 85 changed files with 4,046 additions and 9,800 deletions.
46 changes: 46 additions & 0 deletions backend/query/query_engine_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -737,6 +737,52 @@ TEST_P(QueryEngineTest, ExecuteSqlSelectGetInternalSequenceStateInvalidArg) {
testing::HasSubstr("No matching signature for function")));
}

TEST_P(QueryEngineTest, ExecuteSqlSelectGetTableColumnIdentityState) {
test::ScopedEmulatorFeatureFlagsSetter setter(
{.enable_identity_columns = true});

// When using the postgres dialect, GET_TABLE_COLUMN_IDENTITY_STATE() is
// exposed only via the 'spanner' namespace.
std::string function_call =
"GET_TABLE_COLUMN_IDENTITY_STATE('test_id_table.int64_col')";
if (GetParam() == POSTGRESQL) {
function_call = absl::StrCat("SPANNER.", function_call);
}

// Success case.
std::string sql = absl::StrFormat("SELECT %s AS state", function_call);
ZETASQL_ASSERT_OK_AND_ASSIGN(
QueryResult result,
query_engine().ExecuteSql(Query{sql},
QueryContext{sequence_schema(), reader()}));
ASSERT_NE(result.rows, nullptr);
EXPECT_THAT(GetColumnNames(*result.rows), ElementsAre("state"));
EXPECT_THAT(GetColumnTypes(*result.rows), ElementsAre(Int64Type()));
EXPECT_THAT(GetAllColumnValues(std::move(result.rows)),
IsOkAndHolds(ElementsAre(ElementsAre(NullInt64()))));

// Error case.
std::string invalid_sql =
"SELECT "
"INVALID_SCHEMA.GET_TABLE_COLUMN_IDENTITY_STATE('test_id_"
"table.int64_col') AS state";
if (GetParam() == POSTGRESQL) {
EXPECT_THAT(
query_engine().ExecuteSql(Query{invalid_sql},
QueryContext{sequence_schema(), reader()}),
StatusIs(StatusCode::kNotFound,
HasSubstr("function "
"invalid_schema.get_table_column_identity_state("
"unknown) does not exist")));
} else {
EXPECT_THAT(
query_engine().ExecuteSql(Query{invalid_sql},
QueryContext{sequence_schema(), reader()}),
StatusIs(StatusCode::kInvalidArgument,
HasSubstr("Function not found")));
}
}

TEST_P(QueryEngineTest, ExecuteSqlSelectsOneColumnFromTable) {
ZETASQL_ASSERT_OK_AND_ASSIGN(
QueryResult result,
Expand Down
8 changes: 4 additions & 4 deletions backend/query/query_validator.cc
Original file line number Diff line number Diff line change
Expand Up @@ -575,8 +575,8 @@ absl::Status QueryValidator::ValidateSequenceFunction(
const std::string sequence_name =
node->generic_argument_list(0)->sequence()->sequence()->FullName();
const Sequence* current_sequence =
schema()->FindSequence(sequence_name
);
schema()->FindSequence(sequence_name,
/*exclude_internal=*/true);
if (current_sequence == nullptr) {
return error::SequenceNotFound(sequence_name);
}
Expand All @@ -591,8 +591,8 @@ absl::Status QueryValidator::ValidateSequenceFunction(
node->argument_list(0)->GetAs<zetasql::ResolvedLiteral>()->value();
if (value.type()->IsString()) {
const Sequence* current_sequence =
schema()->FindSequence(value.string_value()
);
schema()->FindSequence(value.string_value(),
/*exclude_internal=*/true);
if (current_sequence == nullptr) {
return error::SequenceNotFound(value.string_value());
}
Expand Down
5 changes: 5 additions & 0 deletions backend/schema/builders/table_builder.h
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,11 @@ class Table::Editor {

const Table* get() const { return instance_; }

Editor& set_name(const std::string& name) {
instance_->name_ = name;
return *this;
}

Editor& add_column(const Column* column) {
instance_->columns_.push_back(column);
instance_->columns_map_[column->Name()] = column;
Expand Down
28 changes: 27 additions & 1 deletion backend/schema/catalog/schema.cc
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include "zetasql/public/types/type.h"
#include "absl/container/flat_hash_map.h"
#include "absl/log/check.h"
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "backend/schema/catalog/change_stream.h"
#include "backend/schema/catalog/check_constraint.h"
Expand Down Expand Up @@ -287,7 +288,28 @@ void DumpColumn(const Column* column, ddl::ColumnDefinition& column_def) {
set_option->set_option_name(ddl::kPGCommitTimestampOptionName);
set_option->set_bool_value(true);
}
if (column->has_default_value()) {
if (column->is_identity_column()) {
ddl::ColumnDefinition::IdentityColumnDefinition* identity_column =
column_def.mutable_identity_column();
ABSL_CHECK_EQ(column->sequences_used().size(), 1); // Crash OK
const Sequence* sequence =
static_cast<const Sequence*>(column->sequences_used().at(0));
if (!sequence->use_default_sequence_kind_option()) {
if (sequence->sequence_kind() == Sequence::BIT_REVERSED_POSITIVE) {
identity_column->set_type(
ddl::ColumnDefinition::IdentityColumnDefinition::
BIT_REVERSED_POSITIVE);
}
}
if (sequence->start_with_counter().has_value()) {
identity_column->set_start_with_counter(
sequence->start_with_counter().value());
}
if (sequence->skip_range_min().has_value()) {
identity_column->set_skip_range_min(sequence->skip_range_min().value());
identity_column->set_skip_range_max(sequence->skip_range_max().value());
}
} else if (column->has_default_value()) {
SetColumnExpression(column, *column_def.mutable_column_default());
}
if (column->is_generated()) {
Expand Down Expand Up @@ -483,6 +505,10 @@ ddl::DDLStatementList Schema::Dump() const {

// Print sequences next, since other schema objects may use them.
for (const Sequence* sequence : sequences_) {
if (sequence->is_internal_use()) {
// Do not print internal sequences.
continue;
}
DumpSequence(sequence,
*ddl_statements.add_statement()->mutable_create_sequence());
}
Expand Down
7 changes: 6 additions & 1 deletion backend/schema/catalog/schema_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1740,7 +1740,12 @@ TEST_F(SchemaTest, PostgreSQLPrintDDLStatementsTestSequences) {
int64_col bigint DEFAULT nextval('myseq'::text) NOT NULL,
string_col character varying,
PRIMARY KEY(int64_col)
))")));
))",
"CREATE TABLE test_id_table (\n"
" int64_col bigint GENERATED BY DEFAULT AS "
"IDENTITY (BIT_REVERSED_POSITIVE) NOT NULL,\n"
" string_col character varying,\n"
" PRIMARY KEY(int64_col)\n)")));
}

TEST_F(SchemaTest, PostgreSQLPrintDDLStatementsTestArrays) {
Expand Down
16 changes: 11 additions & 5 deletions backend/schema/parser/ddl_parser_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -4552,17 +4552,23 @@ TEST(CreateSequence, CanParseCreateSequenceAllOptions) {
)pb")));
}

TEST(CreateSequence, Invalid_NoSequenceKind) {
EXPECT_THAT(ParseDDLStatement(R"sql(
TEST(CreateSequence, AllowNoSequenceKind) {
EXPECT_THAT(
ParseDDLStatement(R"sql(
CREATE SEQUENCE seq OPTIONS (
skip_range_min = 1,
skip_range_max = 1000,
start_with_counter = 1
)
)sql"),
StatusIs(StatusCode::kInvalidArgument,
HasSubstr("CREATE SEQUENCE statements require option "
"`sequence_kind` to be set")));
IsOkAndHolds(test::EqualsProto(R"pb(
create_sequence {
sequence_name: "seq"
set_options { option_name: "skip_range_min" int64_value: 1 }
set_options { option_name: "skip_range_max" int64_value: 1000 }
set_options { option_name: "start_with_counter" int64_value: 1 }
}
)pb")));
}

TEST(CreateSequence, Invalid_EmptyOptionList) {
Expand Down
48 changes: 38 additions & 10 deletions backend/schema/updater/schema_updater.cc
Original file line number Diff line number Diff line change
Expand Up @@ -557,8 +557,9 @@ class SchemaUpdaterImpl {
const Table* table);
absl::Status DropConstraint(const std::string& constraint_name,
const Table* table);
absl::Status AddSynonym(const ddl::AlterTable::AddSynonym& add_synonym,
const Table* table);
absl::Status RenameTo(const ddl::AlterTable::RenameTo& rename_to,
const Table* table);
absl::Status AddSynonym(const std::string& synonym, const Table* table);
absl::Status DropSynonym(const ddl::AlterTable::DropSynonym& drop_synonym,
const Table* table);

Expand Down Expand Up @@ -4071,19 +4072,38 @@ absl::Status SchemaUpdaterImpl::ValidateAlterDatabaseOptions(
return absl::OkStatus();
}

absl::Status SchemaUpdaterImpl::AddSynonym(
const ddl::AlterTable::AddSynonym& add_synonym, const Table* table) {
ZETASQL_RETURN_IF_ERROR(global_names_.AddName("Table", add_synonym.synonym()));
absl::Status SchemaUpdaterImpl::RenameTo(
const ddl::AlterTable::RenameTo& rename_to, const Table* table) {
ZETASQL_RETURN_IF_ERROR(global_names_.AddName("Table", rename_to.name()));
global_names_.RemoveName(table->Name());

const Table* updated_table = nullptr;
ZETASQL_RETURN_IF_ERROR(AlterNode(table, [&](Table::Editor* editor) {
editor->set_synonym(add_synonym.synonym());
editor->set_name(rename_to.name());
updated_table = editor->get();
return absl::OkStatus();
}));
if (SDLObjectName::IsFullyQualifiedName(add_synonym.synonym())) {

// Handle synonyms.
if (!rename_to.synonym().empty()) {
ZETASQL_RETURN_IF_ERROR(AddSynonym(rename_to.synonym(), updated_table));
}

return absl::OkStatus();
}

absl::Status SchemaUpdaterImpl::AddSynonym(const std::string& synonym,
const Table* table) {
ZETASQL_RETURN_IF_ERROR(global_names_.AddName("Table", synonym));
const Table* updated_table = nullptr;
ZETASQL_RETURN_IF_ERROR(AlterNode(table, [&](Table::Editor* editor) {
editor->set_synonym(synonym);
updated_table = editor->get();
return absl::OkStatus();
}));
if (SDLObjectName::IsFullyQualifiedName(synonym)) {
ZETASQL_RETURN_IF_ERROR(AlterInNamedSchema(
add_synonym.synonym(),
[&updated_table](NamedSchema::Editor* editor) -> absl::Status {
synonym, [&updated_table](NamedSchema::Editor* editor) -> absl::Status {
editor->add_synonym(updated_table);
return absl::OkStatus();
}));
Expand Down Expand Up @@ -4495,10 +4515,14 @@ absl::Status SchemaUpdaterImpl::AlterTable(
return error::RowDeletionPolicyDoesNotExist(table->Name());
}
}
case ddl::AlterTable::kRenameTo: {
const auto& rename_to = alter_table.rename_to();
return RenameTo(rename_to, table);
}
case ddl::AlterTable::kAddSynonym: {
const auto& add_synonym = alter_table.add_synonym();
if (table->synonym().empty()) {
return AddSynonym(add_synonym, table);
return AddSynonym(add_synonym.synonym(), table);
} else {
return error::SynonymAlreadyExists(table->synonym(), table->Name());
}
Expand Down Expand Up @@ -5178,6 +5202,10 @@ absl::StatusOr<std::unique_ptr<ddl::DDLStatement>> ParseDDLByDialect(
.enable_nulls_ordering = true,
.enable_generated_column = true,
.enable_column_default = true,
.enable_identity_column =
EmulatorFeatureFlags::instance().flags().enable_identity_columns,
.enable_default_sequence_kind =
EmulatorFeatureFlags::instance().flags().enable_identity_columns,
.enable_jsonb_type = true,
.enable_array_jsonb_type = true,
.enable_create_view = true,
Expand Down
2 changes: 2 additions & 0 deletions backend/schema/updater/schema_updater_tests/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -64,8 +64,10 @@ cc_library(
"check_constraint.cc",
"column_default_values.cc",
"common.cc",
"database_option.cc",
"foreign_key.cc",
"generated_column.cc",
"identity_column.cc",
"index.cc",
"model_test.cc",
"named_schema.cc",
Expand Down
Loading

0 comments on commit 7ad3d59

Please sign in to comment.