From 974a3f9cbd87b8b2f75aa3d3f500b871bde4c4df Mon Sep 17 00:00:00 2001
From: Andrew Farries <andrew.farries@xata.io>
Date: Mon, 14 Aug 2023 09:13:40 +0100
Subject: [PATCH] Add create index operation

---
 pkg/migrations/op_common.go       |  7 ++++
 pkg/migrations/op_create_index.go | 69 +++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+)
 create mode 100644 pkg/migrations/op_create_index.go

diff --git a/pkg/migrations/op_common.go b/pkg/migrations/op_common.go
index 0e38dac38..236dad14e 100644
--- a/pkg/migrations/op_common.go
+++ b/pkg/migrations/op_common.go
@@ -16,6 +16,7 @@ const (
 	OpNameDropTable   OpName = "drop_table"
 	OpNameAddColumn   OpName = "add_column"
 	OpNameDropColumn  OpName = "drop_column"
+	OpNameCreateIndex OpName = "create_index"
 )
 
 func TemporaryName(name string) string {
@@ -86,6 +87,9 @@ func (v *Operations) UnmarshalJSON(data []byte) error {
 		case OpNameDropColumn:
 			item = &OpDropColumn{}
 
+		case OpNameCreateIndex:
+			item = &OpCreateIndex{}
+
 		default:
 			return fmt.Errorf("unknown migration type: %v", opName)
 		}
@@ -133,6 +137,9 @@ func (v Operations) MarshalJSON() ([]byte, error) {
 		case *OpDropColumn:
 			opName = OpNameDropColumn
 
+		case *OpCreateIndex:
+			opName = OpNameCreateIndex
+
 		default:
 			panic(fmt.Errorf("unknown operation for %T", op))
 		}
diff --git a/pkg/migrations/op_create_index.go b/pkg/migrations/op_create_index.go
new file mode 100644
index 000000000..104d6d830
--- /dev/null
+++ b/pkg/migrations/op_create_index.go
@@ -0,0 +1,69 @@
+package migrations
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"strings"
+
+	"github.com/lib/pq"
+
+	"pg-roll/pkg/schema"
+)
+
+type OpCreateIndex struct {
+	Table   string   `json:"table"`
+	Columns []string `json:"columns"`
+}
+
+var _ Operation = (*OpCreateIndex)(nil)
+
+func (o *OpCreateIndex) Start(ctx context.Context, conn *sql.DB, schemaName string, stateSchema string, s *schema.Schema) error {
+	// create index concurrently
+	_, err := conn.ExecContext(ctx, fmt.Sprintf("CREATE INDEX CONCURRENTLY IF NOT EXISTS %s ON %s.%s (%s)",
+		pq.QuoteIdentifier(IndexName(o.Table, o.Columns)),
+		pq.QuoteIdentifier(schemaName),
+		pq.QuoteIdentifier(o.Table),
+		strings.Join(quoteColumnNames(o.Columns), ", ")))
+	return err
+}
+
+func (o *OpCreateIndex) Complete(ctx context.Context, conn *sql.DB) error {
+	// No-op
+	return nil
+}
+
+func (o *OpCreateIndex) Rollback(ctx context.Context, conn *sql.DB) error {
+	// drop the index concurrently
+	_, err := conn.ExecContext(ctx, fmt.Sprintf("DROP INDEX CONCURRENTLY IF EXISTS %s",
+		IndexName(o.Table, o.Columns)))
+
+	return err
+}
+
+func (o *OpCreateIndex) Validate(ctx context.Context, s *schema.Schema) error {
+	table := s.GetTable(o.Table)
+
+	if table == nil {
+		return TableDoesNotExistError{Name: o.Table}
+	}
+
+	for _, column := range o.Columns {
+		if table.GetColumn(column) == nil {
+			return ColumnDoesNotExistError{Table: o.Table, Name: column}
+		}
+	}
+
+	return nil
+}
+
+func IndexName(table string, columns []string) string {
+	return "_pgroll_idx_" + table + "_" + strings.Join(columns, "_")
+}
+
+func quoteColumnNames(columns []string) (quoted []string) {
+	for _, col := range columns {
+		quoted = append(quoted, pq.QuoteIdentifier(col))
+	}
+	return quoted
+}