Merge pull request #54076 from mgartner/backport20.2-53967

release-20.2: sql: include partial index predicates in pg_catalog and pg builtins

pkg/sql/logictest/testdata/logic_test/builtin_function

@@ -2064,13 +2064,33 @@ SELECT pg_catalog.pg_get_indexdef(0)
 NULL
 
 statement ok
-CREATE TABLE test.pg_indexdef_test (a INT, UNIQUE INDEX pg_indexdef_idx (a ASC), INDEX other (a DESC))
+CREATE TYPE testenum AS ENUM ('foo', 'bar', 'baz')
+
+statement ok
+CREATE TABLE test.pg_indexdef_test (
+ a INT,
+ e testenum,
+ UNIQUE INDEX pg_indexdef_idx (a ASC),
+ INDEX pg_indexdef_partial_idx (a) WHERE a > 0,
+ INDEX pg_indexdef_partial_enum_idx (a) WHERE e IN ('foo', 'bar'),
+ INDEX other (a DESC)
+)
 
 query T
 SELECT pg_catalog.pg_get_indexdef((SELECT oid from pg_class WHERE relname='pg_indexdef_idx'))
 ----
 CREATE UNIQUE INDEX pg_indexdef_idx ON test.public.pg_indexdef_test USING btree (a ASC)
 
+query T
+SELECT pg_catalog.pg_get_indexdef((SELECT oid from pg_class WHERE relname='pg_indexdef_partial_idx'))
+----
+CREATE INDEX pg_indexdef_partial_idx ON test.public.pg_indexdef_test USING btree (a ASC) WHERE a > 0
+
+query T
+SELECT pg_catalog.pg_get_indexdef((SELECT oid from pg_class WHERE relname='pg_indexdef_partial_enum_idx'))
+----
+CREATE INDEX pg_indexdef_partial_enum_idx ON test.public.pg_indexdef_test USING btree (a ASC) WHERE e IN ('foo'::public.testenum, 'bar'::public.testenum)
+
 query T
 SELECT pg_catalog.pg_get_indexdef(0, 0, true)
 ----
@@ -2155,6 +2175,8 @@ CREATE TABLE test.pg_constraintdef_test (
  a int,
  b int unique,
  c int check (c > a),
+ d string,
+ UNIQUE INDEX (a) WHERE d = 'foo',
  FOREIGN KEY(a) REFERENCES test.pg_indexdef_test(a) ON DELETE CASCADE
 )
 
@@ -2164,6 +2186,7 @@ FROM pg_catalog.pg_constraint
 WHERE conrelid='pg_constraintdef_test'::regclass
 ----
 UNIQUE (b ASC)
+UNIQUE (a ASC) WHERE d = 'foo'::STRING
 FOREIGN KEY (a) REFERENCES pg_indexdef_test(a) ON DELETE CASCADE
 CHECK ((c > a))
 
@@ -2543,12 +2566,12 @@ NULL
 query I
 SELECT crdb_internal.get_namespace_id(0, 'root_test')
 ----
-61
+63
 
 query I
 SELECT crdb_internal.get_namespace_id(crdb_internal.get_namespace_id(0, 'root_test'), 't')
 ----
-62
+64
 
 query T
 SELECT crdb_internal.get_zone_config(-1)::string
@@ -2598,12 +2621,12 @@ user testuser
 query I
 SELECT crdb_internal.get_namespace_id(0, 'root_test')
 ----
-61
+63
 
 query I
 SELECT crdb_internal.get_namespace_id(crdb_internal.get_namespace_id(0, 'root_test'), 't')
 ----
-62
+64
 
 query T
 SELECT crdb_internal.get_zone_config(crdb_internal.get_namespace_id(0, 'root_test'))::string

pkg/sql/logictest/testdata/logic_test/partial_index

@@ -215,6 +215,21 @@ t10 CREATE TABLE public.t10 (
  FAMILY "primary" (a, b, rowid)
 )
 
+# SHOW CONSTRAINTS includes partial index predicate of UNIQUE partial indexes.
+
+statement ok
+CREATE TABLE t11 (a INT, b INT, UNIQUE INDEX (a) WHERE b > 0)
+
+statement ok
+CREATE UNIQUE INDEX t11_b_key ON t11 (b) WHERE a > 0
+
+query TTTTB colnames
+SHOW CONSTRAINTS FROM t11
+----
+table_name constraint_name constraint_type details validated
+t11 t11_a_key UNIQUE UNIQUE (a ASC) WHERE b > 0 true
+t11 t11_b_key UNIQUE UNIQUE (b ASC) WHERE a > 0 true
+
 # Update a non-indexed column referenced by the predicate.
 
 statement ok

pkg/sql/logictest/testdata/logic_test/pg_catalog

@@ -2535,24 +2535,63 @@ primary CREATE UNIQUE INDEX "primary" ON test.public.geospatial_table USING b
 idxa CREATE INDEX idxa ON test.public.geospatial_table USING gin (a ASC)
 idxb CREATE INDEX idxb ON test.public.geospatial_table USING gin (b ASC)
 
+subtest partial_index
+
+statement ok
+SET DATABASE = test
+
+statement ok
+CREATE TYPE testenum AS ENUM ('foo', 'bar', 'baz')
+
+statement ok
+CREATE TABLE partial_index_table (
+ a INT,
+ b testenum,
+ UNIQUE INDEX (a) WHERE a > 0
+)
+
+statement ok
+CREATE UNIQUE INDEX ON partial_index_table (a) WHERE b IN ('foo', 'bar')
+
+query TT colnames
+SELECT indexname, indexdef
+FROM pg_catalog.pg_indexes
+WHERE tablename = 'partial_index_table'
+----
+indexname indexdef
+primary CREATE UNIQUE INDEX "primary" ON test.public.partial_index_table USING btree (rowid ASC)
+partial_index_table_a_key CREATE UNIQUE INDEX partial_index_table_a_key ON test.public.partial_index_table USING btree (a ASC) WHERE a > 0
+partial_index_table_a_key1 CREATE UNIQUE INDEX partial_index_table_a_key1 ON test.public.partial_index_table USING btree (a ASC) WHERE b IN ('foo'::public.testenum, 'bar'::public.testenum)
+
+query TT colnames
+SELECT conname, condef
+FROM pg_catalog.pg_constraint c JOIN pg_catalog.pg_class t
+ON c.conrelid = t.oid
+WHERE t.relname = 'partial_index_table'
+ORDER BY conname
+----
+conname condef
+partial_index_table_a_key UNIQUE (a ASC) WHERE a > 0
+partial_index_table_a_key1 UNIQUE (a ASC) WHERE b IN ('foo'::public.testenum, 'bar'::public.testenum)
+
 subtest regression_46799
 statement ok
-CREATE TABLE t(x INT DEFAULT 1, y INT DEFAULT 1);
+CREATE TABLE t46799 (x INT DEFAULT 1, y INT DEFAULT 1);
 
 query I
-SELECT adnum FROM pg_attrdef WHERE adrelid = 91
+SELECT adnum FROM pg_attrdef WHERE adrelid = 't46799'::REGCLASS
 ----
 1
 2
 3
 
 statement ok
-ALTER TABLE t DROP COLUMN y;
-ALTER TABLE t ADD COLUMN y INT DEFAULT 1;
+ALTER TABLE t46799 DROP COLUMN y;
+ALTER TABLE t46799 ADD COLUMN y INT DEFAULT 1;
 
 # Make sure after adding and dropping the same column, the adnum for the re-added column increases.
 query I
-select adnum from pg_attrdef WHERE adrelid = 91
+select adnum from pg_attrdef WHERE adrelid = 't46799'::REGCLASS
 ----
 1
 3
@@ -2565,13 +2604,13 @@ CREATE TABLE jt (a INT PRIMARY KEY); INSERT INTO jt VALUES(1); INSERT INTO jt VA
 query ITT
 SELECT a, oid, relname FROM jt INNER LOOKUP JOIN pg_class ON a::oid=oid
 ----
-92 92 jt
+95 95 jt
 
 query ITT
 SELECT a, oid, relname FROM jt LEFT OUTER LOOKUP JOIN pg_class ON a::oid=oid
 ----
 1 NULL NULL
-92 92 jt
+95 95 jt
 
 subtest regression_49207
 statement ok

pkg/sql/pg_catalog.go

@@ -31,6 +31,7 @@ import (
  "github.com/cockroachdb/cockroach/pkg/sql/catalog/schemadesc"
  "github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
  "github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
+ "github.com/cockroachdb/cockroach/pkg/sql/parser"
  "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
  "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
  "github.com/cockroachdb/cockroach/pkg/sql/schemaexpr"
@@ -961,6 +962,14 @@ func populateTableConstraints(
  f.WriteString("UNIQUE (")
  con.Index.ColNamesFormat(f)
  f.WriteByte(')')
+ if con.Index.IsPartial() {
+ pred, err := schemaexpr.FormatExprForDisplay(ctx, table, con.Index.Predicate, p.SemaCtx(), tree.FmtPGCatalog)
+ if err != nil {
+ return err
+ }
+ f.WriteString(" WHERE ")
+ f.WriteString(pred)
+ }
  condef = tree.NewDString(f.CloseAndGetString())
 
  case descpb.ConstraintTypeCheck:
@@ -1203,7 +1212,7 @@ CREATE TABLE pg_catalog.pg_conversion (
  conforencoding INT4,
  contoencoding INT4,
  conproc OID,
- condefault BOOL
+ condefault BOOL
 )`,
  populate: func(ctx context.Context, p *planner, dbContext *dbdesc.Immutable, addRow func(...tree.Datum) error) error {
  return nil
@@ -1826,7 +1835,24 @@ func indexDefFromDescriptor(
  }
  indexDef.Interleave = intlDef
  }
- fmtCtx := tree.NewFmtCtx(tree.FmtPGIndexDef)
+ if index.IsPartial() {
+ // Format the raw predicate for display in order to resolve user-defined
+ // types to a human readable form.
+ //
+ // TODO(mgartner): Avoid parsing the predicate expression twice. It is
+ // parsed in schemaexpr.FormatExprForDisplay and again here.
+ formattedPred, err := schemaexpr.FormatExprForDisplay(ctx, table, index.Predicate, p.SemaCtx(), tree.FmtPGCatalog)
+ if err != nil {
+ return "", err
+ }
+
+ pred, err := parser.ParseExpr(formattedPred)
+ if err != nil {
+ return "", err
+ }
+ indexDef.Predicate = pred
+ }
+ fmtCtx := tree.NewFmtCtx(tree.FmtPGCatalog)
  fmtCtx.FormatNode(&indexDef)
  return fmtCtx.String(), nil
 }

pkg/sql/sem/tree/create.go

@@ -127,7 +127,7 @@ func (node *CreateIndex) Format(ctx *FmtCtx) {
  if node.Unique {
  ctx.WriteString("UNIQUE ")
  }
- if node.Inverted && !ctx.HasFlags(FmtPGIndexDef) {
+ if node.Inverted && !ctx.HasFlags(FmtPGCatalog) {
  ctx.WriteString("INVERTED ")
  }
  ctx.WriteString("INDEX ")
@@ -143,7 +143,7 @@ func (node *CreateIndex) Format(ctx *FmtCtx) {
  }
  ctx.WriteString("ON ")
  ctx.FormatNode(&node.Table)
- if ctx.HasFlags(FmtPGIndexDef) {
+ if ctx.HasFlags(FmtPGCatalog) {
  ctx.WriteString(" USING")
  if node.Inverted {
  ctx.WriteString(" gin")

pkg/sql/sem/tree/format.go

@@ -119,15 +119,12 @@ const (
 
  // FmtPGCatalog is used to produce expressions formatted in a way that's as
  // close as possible to what clients expect to live in pg_catalog (e.g.
- // pg_attrdef.adbin and pg_constraint.condef columns). Specifically, this
- // strips type annotations, since Postgres doesn't know what those are, and
- // adds cast expressions for non-numeric constants.
+ // pg_attrdef.adbin, pg_constraint.condef and pg_indexes.indexdef columns).
+ // Specifically, this strips type annotations (Postgres doesn't know what
+ // those are), adds cast expressions for non-numeric constants, and formats
+ // indexes in Postgres-specific syntax.
  FmtPGCatalog
 
- // FmtPGIndexDef is used to produce CREATE INDEX statements that are
- // compatible with pg_get_indexdef.
- FmtPGIndexDef
-
  // If set, user defined types and datums of user defined types will be
  // formatted in a way that is stable across changes to the underlying type.
  // For type names, this means that they will be formatted as '@id'. For enum