sql/sqlite: support partial indexes

ariga · Mar 8, 2022 · 32981d0 · 32981d0
1 parent 773fc66
commit 32981d0
Show file tree

Hide file tree

Showing 7 changed files with 195 additions and 9 deletions.
diff --git a/internal/integration/testdata/sqlite/index-partial.txt b/internal/integration/testdata/sqlite/index-partial.txt
@@ -0,0 +1,51 @@
+apply 1.hcl
+cmpshow users 1.sql
+
+apply 2.hcl
+cmpshow users 2.sql
+
+-- 1.hcl --
+schema "main" {}
+
+table "users" {
+  schema = schema.main
+  column "name" {
+    null = false
+    type = text
+  }
+  column "active" {
+    null = true
+    type = boolean
+  }
+  index "users_name" {
+    columns = [column.name]
+    where = "active"
+  }
+}
+
+-- 1.sql --
+CREATE TABLE `users` (`name` text NOT NULL, `active` boolean NULL)
+CREATE INDEX `users_name` ON `users` (`name`) WHERE active
+
+-- 2.hcl --
+schema "main" {}
+
+table "users" {
+  schema = schema.main
+  column "name" {
+    null = false
+    type = text
+  }
+  column "active" {
+    null = true
+    type = boolean
+  }
+  index "users_name" {
+    columns = [column.name]
+    where = "active AND name <> ''"
+  }
+}
+
+-- 2.sql --
+CREATE TABLE "users" (`name` text NOT NULL, `active` boolean NULL)
+CREATE INDEX `users_name` ON `users` (`name`) WHERE active AND name <> ''
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
@@ -233,8 +233,8 @@ func fixDefaultQuotes(value schemaspec.Value) error {
 }
 
 // convertIndex converts a sqlspec.Index into a schema.Index.
-func convertIndex(spec *sqlspec.Index, parent *schema.Table) (*schema.Index, error) {
-	idx, err := specutil.Index(spec, parent)
+func convertIndex(spec *sqlspec.Index, t *schema.Table) (*schema.Index, error) {
+	idx, err := specutil.Index(spec, t)
 	if err != nil {
 		return nil, err
 	}
@@ -375,7 +375,7 @@ func indexSpec(idx *schema.Index) (*sqlspec.Index, error) {
 		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("type", strings.ToUpper(i.T)))
 	}
 	if i := (IndexPredicate{}); sqlx.Has(idx.Attrs, &i) && i.P != "" {
-		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("where", i.P))
+		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("where", strconv.Quote(i.P)))
 	}
 	return spec, nil
 }

diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
@@ -289,6 +289,56 @@ table "t" {
 	})
 }
 
+func TestMarshalSpec_IndexPredicate(t *testing.T) {
+	s := &schema.Schema{
+		Name: "test",
+		Tables: []*schema.Table{
+			{
+				Name: "users",
+				Columns: []*schema.Column{
+					{
+						Name: "id",
+						Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int"}},
+					},
+				},
+			},
+		},
+	}
+	s.Tables[0].Schema = s
+	s.Tables[0].Schema = s
+	s.Tables[0].Indexes = []*schema.Index{
+		{
+			Name:   "index",
+			Table:  s.Tables[0],
+			Unique: true,
+			Parts: []*schema.IndexPart{
+				{SeqNo: 0, C: s.Tables[0].Columns[0]},
+			},
+			Attrs: []schema.Attr{
+				&IndexPredicate{P: "id <> 0"},
+			},
+		},
+	}
+	buf, err := MarshalSpec(s, hclState)
+	require.NoError(t, err)
+	const expected = `table "users" {
+  schema = schema.test
+  column "id" {
+    null = false
+    type = int
+  }
+  index "index" {
+    unique  = true
+    columns = [column.id]
+    where   = "id <> 0"
+  }
+}
+schema "test" {
+}
+`
+	require.EqualValues(t, expected, string(buf))
+}
+
 func TestUnmarshalSpec_Identity(t *testing.T) {
 	f := `
 schema "s" {}

diff --git a/sql/sqlite/diff.go b/sql/sqlite/diff.go
@@ -98,10 +98,7 @@ func (d *diff) IsGeneratedIndexName(t *schema.Table, idx *schema.Index) bool {
 // IndexAttrChanged reports if the index attributes were changed.
 func (*diff) IndexAttrChanged(from, to []schema.Attr) bool {
 	var p1, p2 IndexPredicate
-	if sqlx.Has(from, &p1) != sqlx.Has(to, &p2) || p1.P != p2.P {
-		return true
-	}
-	return false
+	return sqlx.Has(from, &p1) != sqlx.Has(to, &p2) || (p1.P != p2.P && p1.P != sqlx.MayWrap(p2.P))
 }
 
 // IndexPartAttrChanged reports if the index-part attributes were changed.

diff --git a/sql/sqlite/diff_test.go b/sql/sqlite/diff_test.go
@@ -165,12 +165,14 @@ func TestDiff_TableDiff(t *testing.T) {
 				{Name: "c2_unique", Unique: true, Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[1]}}},
 				{Name: "c3_predicate", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[1]}}},
 				{Name: "c3_desc", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: to.Columns[1]}}},
+				{Name: "c4_predicate", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[1]}}, Attrs: []schema.Attr{&IndexPredicate{P: "(c4 <> NULL)"}}},
 			}
 			to.Indexes = []*schema.Index{
 				{Name: "c1_index", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[0]}}},
 				{Name: "c3_unique", Unique: true, Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: to.Columns[1]}}},
 				{Name: "c3_predicate", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[1]}}, Attrs: []schema.Attr{&IndexPredicate{P: "c3 <> NULL"}}},
 				{Name: "c3_desc", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, Desc: true, C: to.Columns[1]}}},
+				{Name: "c4_predicate", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[1]}}, Attrs: []schema.Attr{&IndexPredicate{P: "c4 <> NULL"}}},
 			}
 			return testcase{
 				name: "indexes",

diff --git a/sql/sqlite/sqlspec.go b/sql/sqlite/sqlspec.go
@@ -2,6 +2,7 @@ package sqlite
 
 import (
 	"reflect"
+	"strconv"
 
 	"ariga.io/atlas/schema/schemaspec"
 	"ariga.io/atlas/schema/schemaspec/schemahcl"
@@ -25,7 +26,23 @@ func MarshalSpec(v interface{}, marshaler schemaspec.Marshaler) ([]byte, error)
 // ForeignKeySpecs into ForeignKeys, as the target tables do not necessarily exist in the schema
 // at this point. Instead, the linking is done by the convertSchema function.
 func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, error) {
-	return specutil.Table(spec, parent, convertColumn, specutil.PrimaryKey, specutil.Index, specutil.Check)
+	return specutil.Table(spec, parent, convertColumn, specutil.PrimaryKey, convertIndex, specutil.Check)
+}
+
+// convertIndex converts a sqlspec.Index into a schema.Index.
+func convertIndex(spec *sqlspec.Index, t *schema.Table) (*schema.Index, error) {
+	idx, err := specutil.Index(spec, t)
+	if err != nil {
+		return nil, err
+	}
+	if attr, ok := spec.Attr("where"); ok {
+		p, err := attr.String()
+		if err != nil {
+			return nil, err
+		}
+		idx.Attrs = append(idx.Attrs, &IndexPredicate{P: p})
+	}
+	return idx, nil
 }
 
 // convertColumn converts a sqlspec.Column into a schema.Column.
@@ -62,12 +79,23 @@ func tableSpec(tab *schema.Table) (*sqlspec.Table, error) {
 		tab,
 		columnSpec,
 		specutil.FromPrimaryKey,
-		specutil.FromIndex,
+		indexSpec,
 		specutil.FromForeignKey,
 		specutil.FromCheck,
 	)
 }
 
+func indexSpec(idx *schema.Index) (*sqlspec.Index, error) {
+	spec, err := specutil.FromIndex(idx)
+	if err != nil {
+		return nil, err
+	}
+	if i := (IndexPredicate{}); sqlx.Has(idx.Attrs, &i) && i.P != "" {
+		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("where", strconv.Quote(i.P)))
+	}
+	return spec, nil
+}
+
 // columnSpec converts from a concrete SQLite schema.Column into a sqlspec.Column.
 func columnSpec(c *schema.Column, _ *schema.Table) (*sqlspec.Column, error) {
 	s, err := specutil.FromColumn(c, columnTypeSpec)

diff --git a/sql/sqlite/sqlspec_test.go b/sql/sqlite/sqlspec_test.go
@@ -41,6 +41,7 @@ table "table" {
 			table.table.column.id,
 			table.table.column.age,
 		]
+		where = "age <> 0"
 	}
 	foreign_key "accounts" {
 		columns = [
@@ -148,6 +149,9 @@ table "accounts" {
 				{SeqNo: 0, C: exp.Tables[0].Columns[0]},
 				{SeqNo: 1, C: exp.Tables[0].Columns[1]},
 			},
+			Attrs: []schema.Attr{
+				&IndexPredicate{P: "age <> 0"},
+			},
 		},
 	}
 	exp.Tables[0].ForeignKeys = []*schema.ForeignKey{
@@ -208,6 +212,60 @@ schema "test" {
 	require.EqualValues(t, expected, string(buf))
 }
 
+func TestMarshalSpec_IndexPredicate(t *testing.T) {
+	s := &schema.Schema{
+		Name: "test",
+		Tables: []*schema.Table{
+			{
+				Name: "users",
+				Columns: []*schema.Column{
+					{
+						Name: "id",
+						Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int"}},
+						Attrs: []schema.Attr{
+							&AutoIncrement{},
+						},
+					},
+				},
+			},
+		},
+	}
+	s.Tables[0].Schema = s
+	s.Tables[0].Schema = s
+	s.Tables[0].Indexes = []*schema.Index{
+		{
+			Name:   "index",
+			Table:  s.Tables[0],
+			Unique: true,
+			Parts: []*schema.IndexPart{
+				{SeqNo: 0, C: s.Tables[0].Columns[0]},
+			},
+			Attrs: []schema.Attr{
+				&IndexPredicate{P: "id <> 0"},
+			},
+		},
+	}
+	buf, err := MarshalSpec(s, hclState)
+	require.NoError(t, err)
+	const expected = `table "users" {
+  schema = schema.test
+  column "id" {
+    null           = false
+    type           = int
+    auto_increment = true
+  }
+  index "index" {
+    unique  = true
+    columns = [column.id]
+    where   = "id <> 0"
+  }
+}
+schema "test" {
+}
+`
+	require.EqualValues(t, expected, string(buf))
+}
+
 func TestTypes(t *testing.T) {
 	for _, tt := range []struct {
 		typeExpr string