Skip to content

Commit

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

Project import generated by Copybara.
  • Loading branch information
@olavloite
olavloite authored Dec 6, 2024
2 parents 0084b2c + 8482442 commit d3e265a
Show file tree
Hide file tree
Showing 54 changed files with 1,002 additions and 2,731 deletions.
45 changes: 38 additions & 7 deletions backend/schema/parser/ddl_parser.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -815,8 +815,14 @@ void VisitColumnDefaultClauseNode(const SimpleNode* node,
absl::StrCat(ExtractTextForNode(child, ddl_text)));
}

void VisitColumnNode(const SimpleNode* node, ColumnDefinition* column,
absl::string_view ddl_text,
// If containing_table is not null, then we support parsing a PRIMARY KEY clause
// attached to this column; if present, this column's name will be set as the
// sole primary key component in the containing table, or a syntax error will be
// generated if the table's primary key has already been set by a previous
// column. It is expected that higher-level parsing code using this feature will
// also check that there is no table-level PRIMARY KEY clause in that case.
void VisitColumnNode(const SimpleNode* node, absl::string_view ddl_text,
ColumnDefinition* column, CreateTable* containing_table,
std::vector<std::string>* errors) {
CheckNode(node, JJTCOLUMN_DEF);
column->set_column_name(GetChildNode(node, 0, JJTNAME)->image());
Expand All @@ -833,6 +839,18 @@ void VisitColumnNode(const SimpleNode* node, ColumnDefinition* column,
case JJTNOT_NULL:
column->set_not_null(true);
break;
case JJTPRIMARY_KEY:
if (containing_table != nullptr) {
if (containing_table->primary_key().empty()) {
containing_table->add_primary_key()->set_key_name(
column->column_name());
} else {
errors->push_back("Multiple columns declared as PRIMARY KEY");
}
} else {
errors->push_back("Unexpected PRIMARY KEY clause");
}
break;
case JJTGENERATION_CLAUSE:
VisitGenerationClauseNode(child, column, ddl_text);
break;
Expand Down Expand Up @@ -1138,11 +1156,13 @@ void VisitCreateTableNode(const SimpleNode* node, CreateTable* table,
GetQualifiedIdentifier(GetChildNode(node, offset, JJTNAME)));
offset++;

bool has_primary_key = false;

for (int i = offset; i < node->jjtGetNumChildren(); i++) {
SimpleNode* child = GetChildNode(node, i);
switch (child->getId()) {
case JJTCOLUMN_DEF:
VisitColumnNode(child, table->add_column(), ddl_text, errors);
VisitColumnNode(child, ddl_text, table->add_column(), table, errors);
break;
case JJTFOREIGN_KEY:
VisitForeignKeyNode(child, table->add_foreign_key(), errors);
Expand All @@ -1152,7 +1172,12 @@ void VisitCreateTableNode(const SimpleNode* node, CreateTable* table,
errors);
break;
case JJTPRIMARY_KEY:
VisitKeyNode(child, table->mutable_primary_key(), errors);
if (table->primary_key().empty()) {
VisitKeyNode(child, table->mutable_primary_key(), errors);
} else {
errors->push_back("Cannot specify both table and column PRIMARY KEY");
}
has_primary_key = true;
break;
case JJTSYNONYM_CLAUSE:
table->set_synonym(GetChildNode(child, 0)->image());
Expand All @@ -1168,6 +1193,11 @@ void VisitCreateTableNode(const SimpleNode* node, CreateTable* table,
ABSL_LOG(FATAL) << "Unexpected table info: " << child->toString();
}
}

if (table->primary_key().empty() && !has_primary_key) {
// Even for singleton tables, a table-level PRIMARY KEY() must be specified.
errors->push_back("Must specify either table or column PRIMARY KEY");
}
}

void VisitCreateFunctionNode(const SimpleNode* node,
Expand Down Expand Up @@ -2140,9 +2170,9 @@ void VisitAlterTableNode(const SimpleNode* node, absl::string_view ddl_text,
alter_table->mutable_add_column()->set_existence_modifier(
IF_NOT_EXISTS);
}
VisitColumnNode(GetChildNode(node, offset, JJTCOLUMN_DEF),
VisitColumnNode(GetChildNode(node, offset, JJTCOLUMN_DEF), ddl_text,
alter_table->mutable_add_column()->mutable_column(),
ddl_text, errors);
/*containing_table=*/nullptr, errors);
} break;
case JJTDROP_COLUMN: {
SimpleNode* column = GetChildNode(node, 2, JJTCOLUMN_NAME);
Expand Down Expand Up @@ -2417,7 +2447,8 @@ void VisitModelColumnList(const SimpleNode* node,
for (int i = 0; i < node->jjtGetNumChildren(); i++) {
SimpleNode* child = GetChildNode(node, i);
CheckNode(child, JJTCOLUMN_DEF);
VisitColumnNode(child, columns->Add(), ddl_text, errors);
VisitColumnNode(child, ddl_text, columns->Add(),
/*containing_table=*/nullptr, errors);
}
}

Expand Down
3 changes: 2 additions & 1 deletion backend/schema/parser/ddl_parser.jjt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -256,7 +256,7 @@ void create_table_statement() :
[ if_not_exists() ]
qualified_identifier() #name
"(" [ table_element() ( LOOKAHEAD(2) "," table_element() )* [ "," ] ] ")"
primary_key()
[ primary_key() ]
[ LOOKAHEAD(2) "," table_interleave_clause() ]
[ LOOKAHEAD(2) "," row_deletion_policy_clause() ]
}
Expand All @@ -281,6 +281,7 @@ void column_def() :
| identity_column_clause()
]
[ <HIDDEN> #hidden ]
[ <PRIMARY> <KEY> #primary_key ]
[ options_clause() ]
}

Expand Down
102 changes: 99 additions & 3 deletions backend/schema/parser/ddl_parser_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -272,15 +272,81 @@ TEST(ParseCreateTable, CannotParseCreateTableWithoutName) {
}

TEST(ParseCreateTable, CannotParseCreateTableWithoutPrimaryKey) {
EXPECT_THAT(ParseDDLStatement(
R"sql(
EXPECT_THAT(
ParseDDLStatement(
R"sql(
CREATE TABLE Users (
UserId INT64 NOT NULL,
Name STRING(MAX)
)
)sql"),
StatusIs(StatusCode::kInvalidArgument,
HasSubstr("Must specify either table or column PRIMARY KEY")));
}

TEST(ParseCreateTable, CanParseCreateTableWithPrimaryKeyOnAColumn) {
EXPECT_THAT(
ParseDDLStatement(
R"sql(
CREATE TABLE Users (
UserId INT64 NOT NULL PRIMARY KEY,
Name STRING(MAX)
)
)sql"),
IsOkAndHolds(test::EqualsProto(
R"pb(
create_table {
table_name: "Users"
column { column_name: "UserId" type: INT64 not_null: true }
column { column_name: "Name" type: STRING }
primary_key { key_name: "UserId" }
}
)pb")));
}

TEST(ParseCreateTable, CannotParseCreateTablePrimaryKeyOnTwoColumns) {
EXPECT_THAT(ParseDDLStatement(
R"sql(
CREATE TABLE Users (
UserId INT64 NOT NULL PRIMARY KEY,
Name STRING(MAX) PRIMARY KEY
)
)sql"),
StatusIs(StatusCode::kInvalidArgument,
HasSubstr("Expecting 'PRIMARY' but found 'EOF'")));
HasSubstr("Multiple columns declared as PRIMARY KEY")));
}

TEST(ParseCreateTable, CannotParseCreateTablePrimaryKeyInColumnAndInTable) {
EXPECT_THAT(
ParseDDLStatement(
R"sql(
CREATE TABLE Users (
UserId INT64 NOT NULL PRIMARY KEY,
Name STRING(MAX)
) PRIMARY KEY (UserId)
)sql"),
StatusIs(StatusCode::kInvalidArgument,
HasSubstr("Cannot specify both table and column PRIMARY KEY")));
}

TEST(ParseCreateTable, CannotParseAlterTableWithPrimaryKeyInAddColumn) {
EXPECT_THAT(ParseDDLStatement(
R"sql(
ALTER TABLE Users
ADD COLUMN UserId INT64 NOT NULL PRIMARY KEY
)sql"),
StatusIs(StatusCode::kInvalidArgument,
HasSubstr("Unexpected PRIMARY KEY clause")));
}

TEST(ParseCreateTable, CannotParseAlterTableWithPrimaryKeyInAlterColumn) {
EXPECT_THAT(ParseDDLStatement(
R"sql(
ALTER TABLE Users
ALTER COLUMN UserId INT64 NOT NULL PRIMARY KEY
)sql"),
StatusIs(StatusCode::kInvalidArgument,
HasSubstr("Expecting 'EOF' but found 'PRIMARY'")));
}

TEST(ParseCreateTable, CanParseCreateTableWithOnlyAKeyColumn) {
Expand Down Expand Up @@ -5176,6 +5242,36 @@ TEST(ParseCreateModel, ParseCreateModel) {
)sql"),
StatusIs(StatusCode::kInvalidArgument,
HasSubstr("Encountered ',' while parsing: column_type")));

// Model has an unexpected PRIMARY KEY in input column
EXPECT_THAT(ParseDDLStatement(R"sql(
CREATE MODEL m
INPUT (
f1 INT64,
f2 STRING(MAX) PRIMARY KEY
)
OUTPUT (
l1 BOOL,
l2 ARRAY<FLOAT64>
)
)sql"),
StatusIs(StatusCode::kInvalidArgument,
HasSubstr("Unexpected PRIMARY KEY clause")));

// Model has an unexpected PRIMARY KEY in output column
EXPECT_THAT(ParseDDLStatement(R"sql(
CREATE MODEL m
INPUT (
f1 INT64,
f2 STRING(MAX)
)
OUTPUT (
l1 BOOL PRIMARY KEY,
l2 ARRAY<FLOAT64>
)
)sql"),
StatusIs(StatusCode::kInvalidArgument,
HasSubstr("Unexpected PRIMARY KEY clause")));
}

TEST(ParseAlterModel, ParseAlterModel) {
Expand Down
21 changes: 21 additions & 0 deletions backend/schema/updater/schema_updater_tests/table.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,27 @@ TEST_P(SchemaUpdaterTest, CreateTable_SingleKey) {
EXPECT_THAT(col2, testing::Not(IsKeyColumnOf(t, "ASC")));
}

TEST_P(SchemaUpdaterTest, CreateTable_SingleKeyInline) {
ZETASQL_ASSERT_OK_AND_ASSIGN(auto schema, CreateSchema({R"(
CREATE TABLE T(
col1 INT64 PRIMARY KEY,
col2 STRING(MAX)
))"}));

auto t = schema->FindTable("T");
EXPECT_NE(t, nullptr);
EXPECT_EQ(t->columns().size(), 2);
EXPECT_EQ(t->primary_key().size(), 1);

auto col1 = t->columns()[0];
EXPECT_THAT(col1, ColumnIs("col1", types::Int64Type()));
EXPECT_THAT(col1, IsKeyColumnOf(t, "ASC"));

auto col2 = t->columns()[1];
EXPECT_THAT(col2, ColumnIs("col2", types::StringType()));
EXPECT_THAT(col2, testing::Not(IsKeyColumnOf(t, "ASC")));
}

TEST_P(SchemaUpdaterTest, CreateTable_MultiKey) {
std::unique_ptr<const Schema> schema;
// Changing the ordering for key columns is unsupported in PG.
Expand Down
Loading

0 comments on commit d3e265a

Please sign in to comment.