Repository: ariga/atlas
Pull request #1258 (2022-11-03 11:43:17)

sql/postgres: support inspecting index operator classes


Patch:
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index a709f935..2f9f927b 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -318,12 +318,15 @@ func (i *inspect) addIndexes(s *schema.Schema, rows *sql.Rows) error {
 	names := make(map[string]*schema.Index)
 	for rows.Next() {
 		var (
-			uniq, primary, included                       bool
-			table, name, typ                              string
-			desc, nullsfirst, nullslast                   sql.NullBool
-			column, contype, pred, expr, comment, options sql.NullString
+			uniq, primary, included                                           bool
+			table, name, typ                                                  string
+			desc, nullsfirst, nullslast, opcdefault                           sql.NullBool
+			column, contype, pred, expr, comment, options, opcname, opcparams sql.NullString
 		)
-		if err := rows.Scan(&table, &name, &typ, &column, &included, &primary, &uniq, &contype, &pred, &expr, &desc, &nullsfirst, &nullslast, &comment, &options); err != nil {
+		if err := rows.Scan(
+			&table, &name, &typ, &column, &included, &primary, &uniq, &contype, &pred, &expr, &desc,
+			&nullsfirst, &nullslast, &comment, &options, &opcname, &opcdefault, &opcparams,
+		); err != nil {
 			return fmt.Errorf("postgres: scanning indexes for schema %q: %w", s.Name, err)
 		}
 		t, ok := s.Table(table)
@@ -395,10 +398,34 @@ func (i *inspect) addIndexes(s *schema.Schema, rows *sql.Rows) error {
 		default:
 			return fmt.Errorf("postgres: invalid part for index %q", idx.Name)
 		}
+		if err := mayAppendOps(part, opcname.String, opcparams.String, opcdefault.Bool); err != nil {
+			return err
+		}
 	}
 	return nil
 }
 
+// mayAppendOps appends an operator_class attribute to the part in case it is not the default.
+func mayAppendOps(part *schema.IndexPart, name string, params string, defaults bool) error {
+	if name == "" || defaults && params == "" {
+		return nil
+	}
+	op := &IndexOpClass{Name: name, Default: defaults}
+	switch {
+	case len(params) == 0:
+	case len(params) < 2, params[0] != '{', params[len(params)-1] != '}':
+		return fmt.Errorf("postgres: unexpected operator class parameters format: %q", params)
+	default:
+		for _, e := range strings.Split(params[1:len(params)-1], ",") {
+			if kv := strings.Split(strings.TrimSpace(e), "="); len(kv) == 2 {
+				op.Params = append(op.Params, struct{ N, V string }{N: kv[0], V: kv[1]})
+			}
+		}
+	}
+	part.Attrs = append(part.Attrs, op)
+	return nil
+}
+
 // partitions builds the partition each table in the schema.
 func (i *inspect) partitions(s *schema.Schema) error {
 	for _, t := range s.Tables {
@@ -779,8 +806,17 @@ type (
 		Columns []*schema.Column
 	}
 
-	// Concurrently describes the CONCURRENTLY clause to instruct Postgres
-	// to build or drop the index concurrently without blocking the current table.
+	// IndexOpClass describers operator class of the index part.
+	// https://www.postgresql.org/docs/current/indexes-opclass.html.
+	IndexOpClass struct {
+		schema.Attr
+		Name    string                  // Name of the operator class.
+		Default bool                    // If it is the default operator class.
+		Params  []struct{ N, V string } // Optional parameters.
+	}
+
+	// Concurrently describes the CONCURRENTLY clause to instruct Postgres to
+	// build or drop the index concurrently without blocking the current table.
 	// https://www.postgresql.org/docs/current/sql-createindex.html#SQL-CREATEINDEX-CONCURRENTLY
 	Concurrently struct {
 		schema.Attr
@@ -1072,11 +1108,14 @@ SELECT
 	c.contype AS constraint_type,
 	pg_get_expr(idx.indpred, idx.indrelid) AS predicate,
 	pg_get_indexdef(idx.indexrelid, idx.ord, false) AS expression,
-	pg_index_column_has_property(idx.indexrelid, idx.ord, 'desc') AS desc,
+	pg_index_column_has_property(idx.indexrelid, idx.ord, 'desc') AS isdesc,
 	pg_index_column_has_property(idx.indexrelid, idx.ord, 'nulls_first') AS nulls_first,
 	pg_index_column_has_property(idx.indexrelid, idx.ord, 'nulls_last') AS nulls_last,
 	obj_description(i.oid, 'pg_class') AS comment,
-	i.reloptions AS options
+	i.reloptions AS options,
+	op.opcname AS opclass_name,
+	op.opcdefault AS opclass_default,
+	a2.attoptions AS opclass_params
 FROM
 	(
 		select
@@ -1091,6 +1130,8 @@ FROM
 	LEFT JOIN pg_constraint c ON idx.indexrelid = c.conindid
 	LEFT JOIN pg_attribute a ON (a.attrelid, a.attnum) = (idx.indrelid, idx.key)
 	JOIN pg_am am ON am.oid = i.relam
+	LEFT JOIN pg_opclass op ON op.oid = idx.indclass[idx.ord-1]
+	LEFT JOIN pg_attribute a2 ON (a2.attrelid, a2.attnum) = (idx.indexrelid, idx.ord)
 WHERE
 	n.nspname = $1
 	AND t.relname IN (%s)


Test patch:
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index 75e83be3..0bedd59f 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -155,26 +155,28 @@ table_name | column_name |      data_type      | formatted |  is_nullable |
 users      | id          | bigint              | int8      |  NO          |                                 |                          |                64 |                    |             0 |               |                    |                | NO          |                |                    |                  |                     |                       |         | b       |         |         |    20
 users      | c1          | smallint            | int2      |  NO          |                                 |                          |                16 |                    |             0 |               |                    |                | NO          |                |                    |                  |                     |                       |         | b       |         |         |    21
 users      | parent_id   | bigint              | int8      |  YES         |                                 |                          |                64 |                    |             0 |               |                    |                | NO          |                |                    |                  |                     |                       |         | b       |         |         |    22
+users      | ts          | tsvector            | tsvector  |  NO          |                                 |                          |                   |                    |             0 |               |                    |                | NO          |                |                    |                  |                     |                       |         | b       |         |         |    23
 `))
 				m.ExpectQuery(queryIndexes).
 					WithArgs("public", "users").
 					WillReturnRows(sqltest.Rows(`
-   table_name   |    index_name   | index_type  | column_name | included | primary | unique | constraint_type | predicate             |   expression              | desc | nulls_first | nulls_last | comment   | options
-----------------+-----------------+-------------+-------------+----------+---------+--------+-----------------+-----------------------+---------------------------+------+-------------+------------+-----------+-----------
-users           | idx             | hash        |             | f        | f       | f      |                 |                       | "left"((c11)::text, 100)  | t    | t           | f          | boring    |
-users           | idx1            | btree       |             | f        | f       | f      |                 | (id <> NULL::integer) | "left"((c11)::text, 100)  | t    | t           | f          |           |
-users           | t1_c1_key       | btree       | c1          | f        | f       | t      | u               |                       | c1                        | t    | t           | f          |           |
-users           | t1_pkey         | btree       | id          | f        | t       | t      | p               |                       | id                        | t    | f           | f          |           |
-users           | idx4            | btree       | c1          | f        | f       | t      |                 |                       | c1                        | f    | f           | f          |           |
-users           | idx4            | btree       | id          | f        | f       | t      |                 |                       | id                        | f    | f           | t          |           |
-users           | idx5            | btree       | c1          | f        | f       | t      |                 |                       | c1                        | f    | f           | f          |           |
-users           | idx5            | btree       |             | f        | f       | t      |                 |                       | coalesce(parent_id, 0)    | f    | f           | f          |           |
-users           | idx6            | brin        | c1          | f        | f       | t      |                 |                       |                           | f    | f           | f          |           | {autosummarize=true,pages_per_range=2}
-users           | idx2            | btree       |             | f        | f       | f      |                 |                       | ((c * 2))                 | f    | f           | t          |           | 
-users           | idx2            | btree       | c1          | f        | f       | f      |                 |                       | c                         | f    | f           | t          |           | 
-users           | idx2            | btree       | id          | f        | f       | f      |                 |                       | d                         | f    | f           | t          |           | 
-users           | idx2            | btree       | c1          | t        | f       | f      |                 |                       | c                         |      |             |            |           | 
-users           | idx2            | btree       | parent_id   | t        | f       | f      |                 |                       | d                         |      |             |            |           |
+   table_name   |    index_name   | index_type  | column_name | included | primary | unique | constraint_type | predicate             |   expression              | desc | nulls_first | nulls_last | comment   |                 options               |   opclass_name    | opclass_default | opclass_params 
+----------------+-----------------+-------------+-------------+----------+---------+--------+-----------------+-----------------------+---------------------------+------+-------------+------------+-----------+---------------------------------------+-------------------+-----------------+----------------
+users           | idx             | hash        |             | f        | f       | f      |                 |                       | "left"((c11)::text, 100)  | t    | t           | f          | boring    |                                       |     int4_ops      |        t        |
+users           | idx1            | btree       |             | f        | f       | f      |                 | (id <> NULL::integer) | "left"((c11)::text, 100)  | t    | t           | f          |           |                                       |     int4_ops      |        t        |
+users           | t1_c1_key       | btree       | c1          | f        | f       | t      | u               |                       | c1                        | t    | t           | f          |           |                                       |     int4_ops      |        t        |
+users           | t1_pkey         | btree       | id          | f        | t       | t      | p               |                       | id                        | t    | f           | f          |           |                                       |     int4_ops      |        t        |
+users           | idx4            | btree       | c1          | f        | f       | t      |                 |                       | c1                        | f    | f           | f          |           |                                       |     int4_ops      |        t        |
+users           | idx4            | btree       | id          | f        | f       | t      |                 |                       | id                        | f    | f           | t          |           |                                       |     int4_ops      |        t        |
+users           | idx5            | btree       | c1          | f        | f       | t      |                 |                       | c1                        | f    | f           | f          |           |                                       |     int4_ops      |        t        |
+users           | idx5            | btree       |             | f        | f       | t      |                 |                       | coalesce(parent_id, 0)    | f    | f           | f          |           |                                       |     int4_ops      |        t        |
+users           | idx6            | brin        | c1          | f        | f       | t      |                 |                       |                           | f    | f           | f          |           | {autosummarize=true,pages_per_range=2}|     int4_ops      |        t        |
+users           | idx2            | btree       |             | f        | f       | f      |                 |                       | ((c * 2))                 | f    | f           | t          |           |                                       |     int4_ops      |        t        |
+users           | idx2            | btree       | c1          | f        | f       | f      |                 |                       | c                         | f    | f           | t          |           |                                       |     int4_ops      |        t        |
+users           | idx2            | btree       | id          | f        | f       | f      |                 |                       | d                         | f    | f           | t          |           |                                       |     int4_ops      |        t        |
+users           | idx2            | btree       | c1          | t        | f       | f      |                 |                       | c                         |      |             |            |           |                                       |     int4_ops      |        t        |
+users           | idx2            | btree       | parent_id   | t        | f       | f      |                 |                       | d                         |      |             |            |           |                                       |     int4_ops      |        t        |
+users           | tsx             | gist        | ts          | f        | f       | f      |                 |                       | ts                        |      |             |            |           |                                       |     tsvector_ops  |        f        | {siglen=1}
 `))
 				m.noFKs()
 				m.noChecks()
@@ -186,6 +188,7 @@ users           | idx2            | btree       | parent_id   | t        | f
 					{Name: "id", Type: &schema.ColumnType{Raw: "bigint", Type: &schema.IntegerType{T: "bigint"}}},
 					{Name: "c1", Type: &schema.ColumnType{Raw: "smallint", Type: &schema.IntegerType{T: "smallint"}}},
 					{Name: "parent_id", Type: &schema.ColumnType{Raw: "bigint", Null: true, Type: &schema.IntegerType{T: "bigint"}}},
+					{Name: "ts", Type: &schema.ColumnType{Raw: "tsvector", Type: &TextSearchType{T: "tsvector"}}},
 				}
 				indexes := []*schema.Index{
 					{Name: "idx", Table: t, Attrs: []schema.Attr{&IndexType{T: "hash"}, &schema.Comment{Text: "boring"}}, Parts: []*schema.IndexPart{{SeqNo: 1, X: &schema.RawExpr{X: `"left"((c11)::text, 100)`}, Desc: true, Attrs: []schema.Attr{&IndexColumnProperty{NullsFirst: true}}}}},
@@ -194,7 +197,8 @@ users           | idx2            | btree       | parent_id   | t        | f
 					{Name: "idx4", Unique: true, Table: t, Attrs: []schema.Attr{&IndexType{T: "btree"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[1]}, {SeqNo: 2, C: columns[0], Attrs: []schema.Attr{&IndexColumnProperty{NullsLast: true}}}}},
 					{Name: "idx5", Unique: true, Table: t, Attrs: []schema.Attr{&IndexType{T: "btree"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[1]}, {SeqNo: 2, X: &schema.RawExpr{X: `coalesce(parent_id, 0)`}}}},
 					{Name: "idx6", Unique: true, Table: t, Attrs: []schema.Attr{&IndexType{T: "brin"}, &IndexStorageParams{AutoSummarize: true, PagesPerRange: 2}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[1]}}},
-					{Name: "idx2", Unique: false, Table: t, Attrs: []schema.Attr{&IndexType{T: "btree"}, &IndexInclude{Columns: columns[1:]}}, Parts: []*schema.IndexPart{{SeqNo: 1, X: &schema.RawExpr{X: `((c * 2))`}, Attrs: []schema.Attr{&IndexColumnProperty{NullsLast: true}}}, {SeqNo: 2, C: columns[1], Attrs: []schema.Attr{&IndexColumnProperty{NullsLast: true}}}, {SeqNo: 3, C: columns[0], Attrs: []schema.Attr{&IndexColumnProperty{NullsLast: true}}}}},
+					{Name: "idx2", Unique: false, Table: t, Attrs: []schema.Attr{&IndexType{T: "btree"}, &IndexInclude{Columns: columns[1:3]}}, Parts: []*schema.IndexPart{{SeqNo: 1, X: &schema.RawExpr{X: `((c * 2))`}, Attrs: []schema.Attr{&IndexColumnProperty{NullsLast: true}}}, {SeqNo: 2, C: columns[1], Attrs: []schema.Attr{&IndexColumnProperty{NullsLast: true}}}, {SeqNo: 3, C: columns[0], Attrs: []schema.Attr{&IndexColumnProperty{NullsLast: true}}}}},
+					{Name: "tsx", Unique: false, Table: t, Attrs: []schema.Attr{&IndexType{T: "gist"}}, Parts: []*schema.IndexPart{{SeqNo: 1, C: columns[3], Attrs: []schema.Attr{&IndexOpClass{Name: "tsvector_ops", Params: []struct{ N, V string }{{N: "siglen", V: "1"}}}}}}},
 				}
 				pk := &schema.Index{
 					Name:   "t1_pkey",
@@ -204,7 +208,8 @@ users           | idx2            | btree       | parent_id   | t        | f
 					Parts:  []*schema.IndexPart{{SeqNo: 1, C: columns[0], Desc: true}},
 				}
 				columns[0].Indexes = append(columns[0].Indexes, pk, indexes[3], indexes[6])
-				columns[1].Indexes = indexes[2:]
+				columns[1].Indexes = indexes[2:7]
+				columns[3].Indexes = indexes[7:]
 				require.EqualValues(columns, t.Columns)
 				require.EqualValues(indexes, t.Indexes)
 				require.EqualValues(pk, t.PrimaryKey)