From 29b12be439e0e5d43396f220ce32efa44e1d7de4 Mon Sep 17 00:00:00 2001
From: Nathan VanBenschoten <nvanbenschoten@gmail.com>
Date: Tue, 7 Sep 2021 23:40:23 -0400
Subject: [PATCH] sql: implement information_schema._pg_index_position

Needed for #69010.

This commit implements the `information_schema._pg_index_position` builtin
function. Given an index's OID and an underlying-table column number,
`information_schema._pg_index_position` return the column's position in the
index (or NULL if not there).

The function is implemented as a user-defined function in Postgres here:
https://github.com/postgres/postgres/blob/master/src/backend/catalog/information_schema.sql

Release note: None.

Release justification: None, waiting for v22.1.
---
 .../logictest/testdata/logic_test/pg_builtins | 53 +++++++++++++++++++
 pkg/sql/sem/builtins/pg_builtins.go           | 46 ++++++++++++++++
 2 files changed, 99 insertions(+)

diff --git a/pkg/sql/logictest/testdata/logic_test/pg_builtins b/pkg/sql/logictest/testdata/logic_test/pg_builtins
index 4412bd792978..a3e62850ca31 100644
--- a/pkg/sql/logictest/testdata/logic_test/pg_builtins
+++ b/pkg/sql/logictest/testdata/logic_test/pg_builtins
@@ -340,3 +340,56 @@ SELECT pg_catalog.set_config('woo', 'woo', false)
 
 query error configuration setting.*not supported
 SELECT set_config('vacuum_cost_delay', '0', false)
+
+# information_schema._pg_index_position
+
+statement ok
+CREATE TABLE indexed (
+  a INT PRIMARY KEY,
+  b INT,
+  c INT,
+  d INT,
+  INDEX (b, d),
+  INDEX (c, a)
+);
+
+# TODO(nvanbenschoten): this will not be needed after #69909 lands.
+statement ok
+SET experimental_enable_temp_tables = 'on'
+
+statement ok
+CREATE TEMPORARY VIEW indexes AS
+  SELECT i.relname, indkey, indexrelid
+    FROM pg_catalog.pg_index
+    JOIN pg_catalog.pg_class AS t ON indrelid   = t.oid
+    JOIN pg_catalog.pg_class AS i ON indexrelid = i.oid
+   WHERE t.relname = 'indexed'
+ORDER BY i.relname
+
+query TT
+SELECT relname, indkey FROM indexes ORDER BY relname DESC
+----
+primary          1
+indexed_c_a_idx  3 1
+indexed_b_d_idx  2 4
+
+query TTII
+SELECT relname,
+       indkey,
+       generate_series(1, 4) input,
+       information_schema._pg_index_position(indexrelid, generate_series(1, 4))
+FROM indexes
+ORDER BY relname DESC, input
+----
+primary          1    1  1
+primary          1    2  NULL
+primary          1    3  NULL
+primary          1    4  NULL
+indexed_c_a_idx  3 1  1  2
+indexed_c_a_idx  3 1  2  NULL
+indexed_c_a_idx  3 1  3  1
+indexed_c_a_idx  3 1  4  NULL
+indexed_b_d_idx  2 4  1  NULL
+indexed_b_d_idx  2 4  2  1
+indexed_b_d_idx  2 4  3  NULL
+indexed_b_d_idx  2 4  4  2
diff --git a/pkg/sql/sem/builtins/pg_builtins.go b/pkg/sql/sem/builtins/pg_builtins.go
index 8fe9403a77c2..362eeb7d4c95 100644
--- a/pkg/sql/sem/builtins/pg_builtins.go
+++ b/pkg/sql/sem/builtins/pg_builtins.go
@@ -1913,6 +1913,52 @@ SELECT description
 			Info:       "Return size in bytes of the column provided as an argument",
 			Volatility: tree.VolatilityImmutable,
 		}),
+
+	// Given an index's OID and an underlying-table column number,
+	// _pg_index_position return the column's position in the index
+	// (or NULL if not there).
+	//
+	// NOTE: this could be defined as a user-defined function, like
+	// it is in Postgres:
+	// https://github.com/postgres/postgres/blob/master/src/backend/catalog/information_schema.sql
+	//
+	//  CREATE FUNCTION _pg_index_position(oid, smallint) RETURNS int
+	//      LANGUAGE sql STRICT STABLE
+	//  BEGIN ATOMIC
+	//  SELECT (ss.a).n FROM
+	//    (SELECT information_schema._pg_expandarray(indkey) AS a
+	//     FROM pg_catalog.pg_index WHERE indexrelid = $1) ss
+	//    WHERE (ss.a).x = $2;
+	//  END;
+	//
+	"information_schema._pg_index_position": makeBuiltin(defProps(),
+		tree.Overload{
+			Types: tree.ArgTypes{
+				{"oid", types.Oid},
+				{"col", types.Int2},
+			},
+			ReturnType: tree.FixedReturnType(types.Int),
+			Fn: func(ctx *tree.EvalContext, args tree.Datums) (tree.Datum, error) {
+				r, err := ctx.InternalExecutor.QueryRow(
+					ctx.Ctx(), "information_schema._pg_index_position",
+					ctx.Txn,
+					`SELECT (ss.a).n FROM
+					  (SELECT information_schema._pg_expandarray(indkey) AS a
+					   FROM pg_catalog.pg_index WHERE indexrelid = $1) ss
+            WHERE (ss.a).x = $2`,
+					args[0], args[1])
+				if err != nil {
+					return nil, err
+				}
+				if len(r) == 0 {
+					return tree.DNull, nil
+				}
+				return r[0], nil
+			},
+			Info:       notUsableInfo,
+			Volatility: tree.VolatilityStable,
+		},
+	),
 }
 
 func getSessionVar(ctx *tree.EvalContext, settingName string, missingOk bool) (tree.Datum, error) {