sql: use separate lookahead of WITH for BUCKET_COUNT

Fixes cockroachdb#75604 In cockroachdb#75231 we use the `WITH_LA` to resolve conflicts on the `WITH BUCKET_COUNT` clause. However, it somehow makes the random sql generator generate sqls like this: ``` IMPORT INTO FAMILY . FAMILY . CHAR ( ident , ident ) PUBLICATION DATA ( PLACEHOLDER ) WITH OPTIONS ( BUCKET_COUNT = PLACEHOLDER ); ``` which is formatted into: ``` IMPORT INTO "family"."family"."char"(ident, ident) PUBLICATION DATA ('placeholder') WITH bucket_count = 'placeholder' ``` this fails the parser for sure, since it's not a legit statement in crdb. This PR creates a new lookahead token to prevent the sql generator from generating `IMPORT` statement with `BUCKET_COUNT` in it. Release note: None
chengxiong-ruan · Feb 2, 2022 · fb14793 · fb14793
1 parent d869341
commit fb14793
Show file tree

Hide file tree

Showing 8 changed files with 24 additions and 22 deletions.
diff --git a/docs/generated/sql/bnf/alter_primary_key.bnf b/docs/generated/sql/bnf/alter_primary_key.bnf
@@ -1,7 +1,7 @@
 alter_onetable_stmt ::=
-	'ALTER' 'TABLE' table_name 'ALTER' 'PRIMARY' 'KEY' 'USING' 'COLUMNS' '(' index_params ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets
+	'ALTER' 'TABLE' table_name 'ALTER' 'PRIMARY' 'KEY' 'USING' 'COLUMNS' '(' index_params ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets
 	| 'ALTER' 'TABLE' table_name 'ALTER' 'PRIMARY' 'KEY' 'USING' 'COLUMNS' '(' index_params ')' 'USING' 'HASH'
 	| 'ALTER' 'TABLE' table_name 'ALTER' 'PRIMARY' 'KEY' 'USING' 'COLUMNS' '(' index_params ')' 
-	| 'ALTER' 'TABLE' 'IF' 'EXISTS' table_name 'ALTER' 'PRIMARY' 'KEY' 'USING' 'COLUMNS' '(' index_params ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets
+	| 'ALTER' 'TABLE' 'IF' 'EXISTS' table_name 'ALTER' 'PRIMARY' 'KEY' 'USING' 'COLUMNS' '(' index_params ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets
 	| 'ALTER' 'TABLE' 'IF' 'EXISTS' table_name 'ALTER' 'PRIMARY' 'KEY' 'USING' 'COLUMNS' '(' index_params ')' 'USING' 'HASH'
 	| 'ALTER' 'TABLE' 'IF' 'EXISTS' table_name 'ALTER' 'PRIMARY' 'KEY' 'USING' 'COLUMNS' '(' index_params ')' 
diff --git a/docs/generated/sql/bnf/col_qualification.bnf b/docs/generated/sql/bnf/col_qualification.bnf
@@ -5,7 +5,7 @@ col_qualification ::=
 	| 'CONSTRAINT' constraint_name 'UNIQUE'
 	| 'CONSTRAINT' constraint_name 'PRIMARY' 'KEY'
 	| 'CONSTRAINT' constraint_name 'PRIMARY' 'KEY' 'USING' 'HASH'
-	| 'CONSTRAINT' constraint_name 'PRIMARY' 'KEY' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets
+	| 'CONSTRAINT' constraint_name 'PRIMARY' 'KEY' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets
 	| 'CONSTRAINT' constraint_name 'CHECK' '(' a_expr ')'
 	| 'CONSTRAINT' constraint_name 'DEFAULT' b_expr
 	| 'CONSTRAINT' constraint_name 'ON' 'UPDATE' b_expr
@@ -22,7 +22,7 @@ col_qualification ::=
 	| 'UNIQUE'
 	| 'PRIMARY' 'KEY'
 	| 'PRIMARY' 'KEY' 'USING' 'HASH'
-	| 'PRIMARY' 'KEY' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets
+	| 'PRIMARY' 'KEY' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets
 	| 'CHECK' '(' a_expr ')'
 	| 'DEFAULT' b_expr
 	| 'ON' 'UPDATE' b_expr

diff --git a/docs/generated/sql/bnf/create_index_stmt.bnf b/docs/generated/sql/bnf/create_index_stmt.bnf
@@ -1,5 +1,5 @@
 create_index_stmt ::=
-	'CREATE' ( 'UNIQUE' |  ) 'INDEX' ( 'CONCURRENTLY' |  ) opt_index_name 'ON' table_name ( 'USING' name |  ) '(' ( ( ( func_expr_windowless index_elem_options | '(' a_expr ')' index_elem_options | name index_elem_options ) ) ( ( ',' ( func_expr_windowless index_elem_options | '(' a_expr ')' index_elem_options | name index_elem_options ) ) )* ) ')' ( 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets | 'USING' 'HASH' |  ) ( ( 'COVERING' | 'STORING' | 'INCLUDE' ) '(' name_list ')' |  ) opt_partition_by_index ( 'WITH' '(' ( ( storage_parameter ) ( ( ',' storage_parameter ) )* ) ')' ) opt_where_clause
-	| 'CREATE' ( 'UNIQUE' |  ) 'INDEX' ( 'CONCURRENTLY' |  ) 'IF' 'NOT' 'EXISTS' index_name 'ON' table_name ( 'USING' name |  ) '(' ( ( ( func_expr_windowless index_elem_options | '(' a_expr ')' index_elem_options | name index_elem_options ) ) ( ( ',' ( func_expr_windowless index_elem_options | '(' a_expr ')' index_elem_options | name index_elem_options ) ) )* ) ')' ( 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets | 'USING' 'HASH' |  ) ( ( 'COVERING' | 'STORING' | 'INCLUDE' ) '(' name_list ')' |  ) opt_partition_by_index ( 'WITH' '(' ( ( storage_parameter ) ( ( ',' storage_parameter ) )* ) ')' ) opt_where_clause
+	'CREATE' ( 'UNIQUE' |  ) 'INDEX' ( 'CONCURRENTLY' |  ) opt_index_name 'ON' table_name ( 'USING' name |  ) '(' ( ( ( func_expr_windowless index_elem_options | '(' a_expr ')' index_elem_options | name index_elem_options ) ) ( ( ',' ( func_expr_windowless index_elem_options | '(' a_expr ')' index_elem_options | name index_elem_options ) ) )* ) ')' ( 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets | 'USING' 'HASH' |  ) ( ( 'COVERING' | 'STORING' | 'INCLUDE' ) '(' name_list ')' |  ) opt_partition_by_index ( 'WITH' '(' ( ( storage_parameter ) ( ( ',' storage_parameter ) )* ) ')' ) opt_where_clause
+	| 'CREATE' ( 'UNIQUE' |  ) 'INDEX' ( 'CONCURRENTLY' |  ) 'IF' 'NOT' 'EXISTS' index_name 'ON' table_name ( 'USING' name |  ) '(' ( ( ( func_expr_windowless index_elem_options | '(' a_expr ')' index_elem_options | name index_elem_options ) ) ( ( ',' ( func_expr_windowless index_elem_options | '(' a_expr ')' index_elem_options | name index_elem_options ) ) )* ) ')' ( 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets | 'USING' 'HASH' |  ) ( ( 'COVERING' | 'STORING' | 'INCLUDE' ) '(' name_list ')' |  ) opt_partition_by_index ( 'WITH' '(' ( ( storage_parameter ) ( ( ',' storage_parameter ) )* ) ')' ) opt_where_clause
 
 
diff --git a/docs/generated/sql/bnf/index_def.bnf b/docs/generated/sql/bnf/index_def.bnf
@@ -1,8 +1,8 @@
 index_def ::=
-	'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets 'COVERING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
-	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets 'STORING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
-	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets 'INCLUDE' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
-	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets  opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
+	'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets 'COVERING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
+	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets 'STORING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
+	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets 'INCLUDE' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
+	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets  opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
 	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'COVERING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
 	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'STORING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
 	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'INCLUDE' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
@@ -11,10 +11,10 @@ index_def ::=
 	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')'  'STORING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
 	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')'  'INCLUDE' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
 	| 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')'   opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
-	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets 'COVERING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
-	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets 'STORING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
-	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets 'INCLUDE' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
-	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets  opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
+	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets 'COVERING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
+	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets 'STORING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
+	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets 'INCLUDE' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
+	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets  opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
 	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'COVERING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
 	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'STORING' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause
 	| 'UNIQUE' 'INDEX' opt_index_name '(' index_elem ( ( ',' index_elem ) )* ')' 'USING' 'HASH' 'INCLUDE' '(' name_list ')' opt_partition_by_index opt_with_storage_parameter_list opt_where_clause

diff --git a/docs/generated/sql/bnf/stmt_block.bnf b/docs/generated/sql/bnf/stmt_block.bnf
@@ -1978,7 +1978,7 @@ index_params ::=
 	( index_elem ) ( ( ',' index_elem ) )*
 
 opt_hash_sharded ::=
-	'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' a_expr
+	'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' a_expr
 	| 'USING' 'HASH'
 	| 
 
@@ -3317,7 +3317,7 @@ col_qualification_elem ::=
 	| 'UNIQUE'
 	| 'PRIMARY' 'KEY'
 	| 'PRIMARY' 'KEY' 'USING' 'HASH'
-	| 'PRIMARY' 'KEY' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' a_expr
+	| 'PRIMARY' 'KEY' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' a_expr
 	| 'CHECK' '(' a_expr ')'
 	| 'DEFAULT' b_expr
 	| 'ON' 'UPDATE' b_expr

diff --git a/docs/generated/sql/bnf/table_constraint.bnf b/docs/generated/sql/bnf/table_constraint.bnf
@@ -4,7 +4,7 @@ table_constraint ::=
 	| 'CONSTRAINT' constraint_name 'UNIQUE' '(' index_params ')' 'STORING' '(' name_list ')' opt_partition_by_index opt_where_clause
 	| 'CONSTRAINT' constraint_name 'UNIQUE' '(' index_params ')' 'INCLUDE' '(' name_list ')' opt_partition_by_index opt_where_clause
 	| 'CONSTRAINT' constraint_name 'UNIQUE' '(' index_params ')'  opt_partition_by_index opt_where_clause
-	| 'CONSTRAINT' constraint_name 'PRIMARY' 'KEY' '(' index_params ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets
+	| 'CONSTRAINT' constraint_name 'PRIMARY' 'KEY' '(' index_params ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets
 	| 'CONSTRAINT' constraint_name 'PRIMARY' 'KEY' '(' index_params ')' 'USING' 'HASH'
 	| 'CONSTRAINT' constraint_name 'PRIMARY' 'KEY' '(' index_params ')' 
 	| 'CONSTRAINT' constraint_name 'FOREIGN' 'KEY' '(' name_list ')' 'REFERENCES' table_name opt_column_list key_match reference_actions
@@ -13,7 +13,7 @@ table_constraint ::=
 	| 'UNIQUE' '(' index_params ')' 'STORING' '(' name_list ')' opt_partition_by_index opt_where_clause
 	| 'UNIQUE' '(' index_params ')' 'INCLUDE' '(' name_list ')' opt_partition_by_index opt_where_clause
 	| 'UNIQUE' '(' index_params ')'  opt_partition_by_index opt_where_clause
-	| 'PRIMARY' 'KEY' '(' index_params ')' 'USING' 'HASH' 'WITH' 'BUCKET_COUNT' '=' n_buckets
+	| 'PRIMARY' 'KEY' '(' index_params ')' 'USING' 'HASH' 'WITH_BUCKET_COUNT' 'BUCKET_COUNT' '=' n_buckets
 	| 'PRIMARY' 'KEY' '(' index_params ')' 'USING' 'HASH'
 	| 'PRIMARY' 'KEY' '(' index_params ')' 
 	| 'FOREIGN' 'KEY' '(' name_list ')' 'REFERENCES' table_name opt_column_list key_match reference_actions
diff --git a/pkg/sql/parser/lexer.go b/pkg/sql/parser/lexer.go
@@ -112,8 +112,10 @@ func (l *lexer) Lex(lval *sqlSymType) int {
 
 		case WITH:
 			switch nextID {
-			case TIME, ORDINALITY, BUCKET_COUNT:
+			case TIME, ORDINALITY:
 				lval.id = WITH_LA
+			case BUCKET_COUNT:
+				lval.id = WITH_BUCKET_COUNT_LA
 			}
 		case NULLS:
 			switch nextID {

diff --git a/pkg/sql/parser/sql.y b/pkg/sql/parser/sql.y
@@ -883,7 +883,7 @@ func (u *sqlSymUnion) setVar() *tree.SetVar {
 // to differentiate `RESET var` from `RESET ALL`. ROLE_ALL and USER_ALL are
 // used in ALTER ROLE statements that affect all roles.
 %token NOT_LA NULLS_LA WITH_LA AS_LA GENERATED_ALWAYS GENERATED_BY_DEFAULT RESET_ALL ROLE_ALL
-%token USER_ALL ON_LA
+%token USER_ALL ON_LA WITH_BUCKET_COUNT_LA
 
 %union {
   id    int32
@@ -6995,7 +6995,7 @@ col_qualification_elem:
     ShardBuckets: tree.DefaultVal{},
   }
 }
-| PRIMARY KEY USING HASH WITH_LA BUCKET_COUNT '=' a_expr
+| PRIMARY KEY USING HASH WITH_BUCKET_COUNT_LA BUCKET_COUNT '=' a_expr
 {
   $$.val = tree.ShardedPrimaryKeyConstraint{
     Sharded: true,
@@ -7323,7 +7323,7 @@ opt_storing:
   }
 
 opt_hash_sharded:
-  USING HASH WITH_LA BUCKET_COUNT '=' a_expr
+  USING HASH WITH_BUCKET_COUNT_LA BUCKET_COUNT '=' a_expr
   {
     $$.val = &tree.ShardedIndexDef{
       ShardBuckets: $6.expr(),