From 3a0c88c46d037a11abaa2e4c3ab13356db3ce1f0 Mon Sep 17 00:00:00 2001
From: Eric Wang <ericw@cockroachlabs.com>
Date: Wed, 7 Apr 2021 23:05:01 -0400
Subject: [PATCH] sql: wrap typed Array within IndirectionExpr in ParenExpr

This patch adds a case that https://github.com/cockroachdb/cockroach/pull/63103 missed.

When we have something like `ARRAY['a'::typ, 'b'::typ]`,
we format it to `ARRAY['a'::typ, 'b'::typ]:::public.typ[]`.
This causes issues if used in an `IndirectionExpr`, such as
`ARRAY['a'::typ, 'b'::typ]:::public.typ[][1]`, as the `[1]`
is interpreted as part of the type.
This would result in errors when trying to use these
expressions in default expressions.
This patch checks for this specific case, and instead
formats it as `(ARRAY['a'::typ, 'b'::typ]:::public.typ[])[1]`.

Release note: None
---
 pkg/sql/logictest/testdata/logic_test/enums | 22 +++++++++++++++++++++
 pkg/sql/parser/testdata/parse/create_table  |  8 ++++++++
 pkg/sql/sem/tree/expr.go                    | 14 ++++++++++---
 pkg/sql/sem/tree/type_check_test.go         |  1 +
 4 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/pkg/sql/logictest/testdata/logic_test/enums b/pkg/sql/logictest/testdata/logic_test/enums
index caf1fb5d0023..e4d7971ab486 100644
--- a/pkg/sql/logictest/testdata/logic_test/enums
+++ b/pkg/sql/logictest/testdata/logic_test/enums
@@ -1363,3 +1363,25 @@ query T
 SELECT * FROM arr_t3
 ----
 c
+
+statement ok
+CREATE TABLE arr_t4 (i typ DEFAULT ARRAY['a'::typ][1], j typ DEFAULT ARRAY['a'::typ, 'b'::typ, 'c'::typ][2])
+
+statement ok
+INSERT INTO arr_t4 VALUES (default, default)
+
+query TT
+SELECT * from arr_t4
+----
+a  b
+
+statement ok
+CREATE TABLE arr_t5 (i typ[] default ARRAY['a'::typ, 'b'::typ])
+
+statement ok
+INSERT INTO arr_t5 VALUES (default)
+
+query T
+SELECT * from arr_t5
+----
+{a,b}
diff --git a/pkg/sql/parser/testdata/parse/create_table b/pkg/sql/parser/testdata/parse/create_table
index d88d57abe092..55cb53e00740 100644
--- a/pkg/sql/parser/testdata/parse/create_table
+++ b/pkg/sql/parser/testdata/parse/create_table
@@ -2150,3 +2150,11 @@ CREATE TABLE arr_t (i INT8 DEFAULT ((('{' || '1') || '}')::INT8[])[1]) -- normal
 CREATE TABLE arr_t (i INT8 DEFAULT (((((((((('{') || ('1'))) || ('}'))))::INT8[])))[(1)])) -- fully parenthetized
 CREATE TABLE arr_t (i INT8 DEFAULT (((_ || _) || _)::INT8[])[_]) -- literals removed
 CREATE TABLE _ (_ INT8 DEFAULT ((('{' || '1') || '}')::INT8[])[1]) -- identifiers removed
+
+parse
+CREATE TABLE arr_t (i INT8 DEFAULT (ARRAY[1, 2, 3]::int[])[2])
+----
+CREATE TABLE arr_t (i INT8 DEFAULT (ARRAY[1, 2, 3]::INT8[])[2]) -- normalized!
+CREATE TABLE arr_t (i INT8 DEFAULT (((((ARRAY[(1), (2), (3)])::INT8[])))[(2)])) -- fully parenthetized
+CREATE TABLE arr_t (i INT8 DEFAULT (ARRAY[_, _, __more1__]::INT8[])[_]) -- literals removed
+CREATE TABLE _ (_ INT8 DEFAULT (ARRAY[1, 2, 3]::INT8[])[2]) -- identifiers removed
diff --git a/pkg/sql/sem/tree/expr.go b/pkg/sql/sem/tree/expr.go
index be4e622dd84f..d5c6b4ddac20 100644
--- a/pkg/sql/sem/tree/expr.go
+++ b/pkg/sql/sem/tree/expr.go
@@ -1624,10 +1624,18 @@ type IndirectionExpr struct {
 
 // Format implements the NodeFormatter interface.
 func (node *IndirectionExpr) Format(ctx *FmtCtx) {
-	// If the subExpr is a CastExpr, we need to wrap it in a ParenExpr,
-	// otherwise the indirection will get interpreted as part of the type.
+	// If the sub expression is a CastExpr or an Array that has a type,
+	// we need to wrap it in a ParenExpr, otherwise the indirection
+	// will get interpreted as part of the type.
 	// Ex. ('{a}'::_typ)[1] vs. '{a}'::_typ[1].
-	if _, isCast := node.Expr.(*CastExpr); isCast {
+	// Ex. (ARRAY['a'::typ]:::typ[])[1] vs. ARRAY['a'::typ]:::typ[][1].
+	var annotateArray bool
+	if arr, ok := node.Expr.(*Array); ctx.HasFlags(FmtParsable) && ok && arr.typ != nil {
+		if arr.typ.ArrayContents().Family() != types.UnknownFamily {
+			annotateArray = true
+		}
+	}
+	if _, isCast := node.Expr.(*CastExpr); isCast || annotateArray {
 		withParens := ParenExpr{Expr: node.Expr}
 		exprFmtWithParen(ctx, &withParens)
 	} else {
diff --git a/pkg/sql/sem/tree/type_check_test.go b/pkg/sql/sem/tree/type_check_test.go
index edde9bb9d52f..80665999e013 100644
--- a/pkg/sql/sem/tree/type_check_test.go
+++ b/pkg/sql/sem/tree/type_check_test.go
@@ -199,6 +199,7 @@ func TestTypeCheck(t *testing.T) {
 
 		{`(('{' || 'a' ||'}')::STRING[])[1]::STRING`, `((('{':::STRING || 'a':::STRING) || '}':::STRING)::STRING[])[1:::INT8]`},
 		{`(('{' || '1' ||'}')::INT[])[1]`, `((('{':::STRING || '1':::STRING) || '}':::STRING)::INT8[])[1:::INT8]`},
+		{`(ARRAY[1, 2, 3]::int[])[2]`, `(ARRAY[1:::INT8, 2:::INT8, 3:::INT8]:::INT8[])[2:::INT8]`},
 	}
 	ctx := context.Background()
 	for _, d := range testData {