From 8156a5d203dfced04c9a856fc28df63716c23ba7 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/sem/tree/expr.go                    | 14 ++++++++++---
 pkg/sql/sem/tree/type_check_test.go         |  1 +
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/pkg/sql/logictest/testdata/logic_test/enums b/pkg/sql/logictest/testdata/logic_test/enums
index 62b097e6a287..38ec154dc358 100644
--- a/pkg/sql/logictest/testdata/logic_test/enums
+++ b/pkg/sql/logictest/testdata/logic_test/enums
@@ -1218,3 +1218,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/sem/tree/expr.go b/pkg/sql/sem/tree/expr.go
index 30cb83da35cb..8bb352c8ba4f 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 862d72d1606b..158c1272806f 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 {