Skip to content

Commit 32981d0

Browse files
a8ma8m
committed
sql/sqlite: support partial indexes
1 parent 773fc66 commit 32981d0

File tree

7 files changed

+195
-9
lines changed

7 files changed

+195
-9
lines changed

‎internal/integration/testdata/sqlite/index-partial.txt

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
apply 1.hcl
2+
cmpshow users 1.sql
3+
4+
apply 2.hcl
5+
cmpshow users 2.sql
6+
7+
-- 1.hcl --
8+
schema "main" {}
9+
10+
table "users" {
11+
schema = schema.main
12+
column "name" {
13+
null = false
14+
type = text
15+
}
16+
column "active" {
17+
null = true
18+
type = boolean
19+
}
20+
index "users_name" {
21+
columns = [column.name]
22+
where = "active"
23+
}
24+
}
25+
26+
-- 1.sql --
27+
CREATE TABLE `users` (`name` text NOT NULL, `active` boolean NULL)
28+
CREATE INDEX `users_name` ON `users` (`name`) WHERE active
29+
30+
-- 2.hcl --
31+
schema "main" {}
32+
33+
table "users" {
34+
schema = schema.main
35+
column "name" {
36+
null = false
37+
type = text
38+
}
39+
column "active" {
40+
null = true
41+
type = boolean
42+
}
43+
index "users_name" {
44+
columns = [column.name]
45+
where = "active AND name <> ''"
46+
}
47+
}
48+
49+
-- 2.sql --
50+
CREATE TABLE "users" (`name` text NOT NULL, `active` boolean NULL)
51+
CREATE INDEX `users_name` ON `users` (`name`) WHERE active AND name <> ''

‎sql/postgres/sqlspec.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -233,8 +233,8 @@ func fixDefaultQuotes(value schemaspec.Value) error {
233233
}
234234

235235
// convertIndex converts a sqlspec.Index into a schema.Index.
236-
func convertIndex(spec *sqlspec.Index, parent *schema.Table) (*schema.Index, error) {
237-
idx, err := specutil.Index(spec, parent)
236+
func convertIndex(spec *sqlspec.Index, t *schema.Table) (*schema.Index, error) {
237+
idx, err := specutil.Index(spec, t)
238238
if err != nil {
239239
return nil, err
240240
}
@@ -375,7 +375,7 @@ func indexSpec(idx *schema.Index) (*sqlspec.Index, error) {
375375
spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("type", strings.ToUpper(i.T)))
376376
}
377377
if i := (IndexPredicate{}); sqlx.Has(idx.Attrs, &i) && i.P != "" {
378-
spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("where", i.P))
378+
spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("where", strconv.Quote(i.P)))
379379
}
380380
return spec, nil
381381
}

‎sql/postgres/sqlspec_test.go

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,56 @@ table "t" {
289289
})
290290
}
291291

292+
func TestMarshalSpec_IndexPredicate(t *testing.T) {
293+
s := &schema.Schema{
294+
Name: "test",
295+
Tables: []*schema.Table{
296+
{
297+
Name: "users",
298+
Columns: []*schema.Column{
299+
{
300+
Name: "id",
301+
Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int"}},
302+
},
303+
},
304+
},
305+
},
306+
}
307+
s.Tables[0].Schema = s
308+
s.Tables[0].Schema = s
309+
s.Tables[0].Indexes = []*schema.Index{
310+
{
311+
Name: "index",
312+
Table: s.Tables[0],
313+
Unique: true,
314+
Parts: []*schema.IndexPart{
315+
{SeqNo: 0, C: s.Tables[0].Columns[0]},
316+
},
317+
Attrs: []schema.Attr{
318+
&IndexPredicate{P: "id <> 0"},
319+
},
320+
},
321+
}
322+
buf, err := MarshalSpec(s, hclState)
323+
require.NoError(t, err)
324+
const expected = `table "users" {
325+
schema = schema.test
326+
column "id" {
327+
null = false
328+
type = int
329+
}
330+
index "index" {
331+
unique = true
332+
columns = [column.id]
333+
where = "id <> 0"
334+
}
335+
}
336+
schema "test" {
337+
}
338+
`
339+
require.EqualValues(t, expected, string(buf))
340+
}
341+
292342
func TestUnmarshalSpec_Identity(t *testing.T) {
293343
f := `
294344
schema "s" {}

‎sql/sqlite/diff.go

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -98,10 +98,7 @@ func (d *diff) IsGeneratedIndexName(t *schema.Table, idx *schema.Index) bool {
9898
// IndexAttrChanged reports if the index attributes were changed.
9999
func (*diff) IndexAttrChanged(from, to []schema.Attr) bool {
100100
var p1, p2 IndexPredicate
101-
if sqlx.Has(from, &p1) != sqlx.Has(to, &p2) || p1.P != p2.P {
102-
return true
103-
}
104-
return false
101+
return sqlx.Has(from, &p1) != sqlx.Has(to, &p2) || (p1.P != p2.P && p1.P != sqlx.MayWrap(p2.P))
105102
}
106103

107104
// IndexPartAttrChanged reports if the index-part attributes were changed.

‎sql/sqlite/diff_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,12 +165,14 @@ func TestDiff_TableDiff(t *testing.T) {
165165
{Name: "c2_unique", Unique: true, Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[1]}}},
166166
{Name: "c3_predicate", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[1]}}},
167167
{Name: "c3_desc", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: to.Columns[1]}}},
168+
{Name: "c4_predicate", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[1]}}, Attrs: []schema.Attr{&IndexPredicate{P: "(c4 <> NULL)"}}},
168169
}
169170
to.Indexes = []*schema.Index{
170171
{Name: "c1_index", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[0]}}},
171172
{Name: "c3_unique", Unique: true, Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: to.Columns[1]}}},
172173
{Name: "c3_predicate", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[1]}}, Attrs: []schema.Attr{&IndexPredicate{P: "c3 <> NULL"}}},
173174
{Name: "c3_desc", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, Desc: true, C: to.Columns[1]}}},
175+
{Name: "c4_predicate", Table: from, Parts: []*schema.IndexPart{{SeqNo: 1, C: from.Columns[1]}}, Attrs: []schema.Attr{&IndexPredicate{P: "c4 <> NULL"}}},
174176
}
175177
return testcase{
176178
name: "indexes",

‎sql/sqlite/sqlspec.go

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package sqlite
22

33
import (
44
"reflect"
5+
"strconv"
56

67
"ariga.io/atlas/schema/schemaspec"
78
"ariga.io/atlas/schema/schemaspec/schemahcl"
@@ -25,7 +26,23 @@ func MarshalSpec(v interface{}, marshaler schemaspec.Marshaler) ([]byte, error)
2526
// ForeignKeySpecs into ForeignKeys, as the target tables do not necessarily exist in the schema
2627
// at this point. Instead, the linking is done by the convertSchema function.
2728
func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, error) {
28-
return specutil.Table(spec, parent, convertColumn, specutil.PrimaryKey, specutil.Index, specutil.Check)
29+
return specutil.Table(spec, parent, convertColumn, specutil.PrimaryKey, convertIndex, specutil.Check)
30+
}
31+
32+
// convertIndex converts a sqlspec.Index into a schema.Index.
33+
func convertIndex(spec *sqlspec.Index, t *schema.Table) (*schema.Index, error) {
34+
idx, err := specutil.Index(spec, t)
35+
if err != nil {
36+
return nil, err
37+
}
38+
if attr, ok := spec.Attr("where"); ok {
39+
p, err := attr.String()
40+
if err != nil {
41+
return nil, err
42+
}
43+
idx.Attrs = append(idx.Attrs, &IndexPredicate{P: p})
44+
}
45+
return idx, nil
2946
}
3047

3148
// convertColumn converts a sqlspec.Column into a schema.Column.
@@ -62,12 +79,23 @@ func tableSpec(tab *schema.Table) (*sqlspec.Table, error) {
6279
tab,
6380
columnSpec,
6481
specutil.FromPrimaryKey,
65-
specutil.FromIndex,
82+
indexSpec,
6683
specutil.FromForeignKey,
6784
specutil.FromCheck,
6885
)
6986
}
7087

88+
func indexSpec(idx *schema.Index) (*sqlspec.Index, error) {
89+
spec, err := specutil.FromIndex(idx)
90+
if err != nil {
91+
return nil, err
92+
}
93+
if i := (IndexPredicate{}); sqlx.Has(idx.Attrs, &i) && i.P != "" {
94+
spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("where", strconv.Quote(i.P)))
95+
}
96+
return spec, nil
97+
}
98+
7199
// columnSpec converts from a concrete SQLite schema.Column into a sqlspec.Column.
72100
func columnSpec(c *schema.Column, _ *schema.Table) (*sqlspec.Column, error) {
73101
s, err := specutil.FromColumn(c, columnTypeSpec)

‎sql/sqlite/sqlspec_test.go

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ table "table" {
4141
table.table.column.id,
4242
table.table.column.age,
4343
]
44+
where = "age <> 0"
4445
}
4546
foreign_key "accounts" {
4647
columns = [
@@ -148,6 +149,9 @@ table "accounts" {
148149
{SeqNo: 0, C: exp.Tables[0].Columns[0]},
149150
{SeqNo: 1, C: exp.Tables[0].Columns[1]},
150151
},
152+
Attrs: []schema.Attr{
153+
&IndexPredicate{P: "age <> 0"},
154+
},
151155
},
152156
}
153157
exp.Tables[0].ForeignKeys = []*schema.ForeignKey{
@@ -208,6 +212,60 @@ schema "test" {
208212
require.EqualValues(t, expected, string(buf))
209213
}
210214

215+
func TestMarshalSpec_IndexPredicate(t *testing.T) {
216+
s := &schema.Schema{
217+
Name: "test",
218+
Tables: []*schema.Table{
219+
{
220+
Name: "users",
221+
Columns: []*schema.Column{
222+
{
223+
Name: "id",
224+
Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int"}},
225+
Attrs: []schema.Attr{
226+
&AutoIncrement{},
227+
},
228+
},
229+
},
230+
},
231+
},
232+
}
233+
s.Tables[0].Schema = s
234+
s.Tables[0].Schema = s
235+
s.Tables[0].Indexes = []*schema.Index{
236+
{
237+
Name: "index",
238+
Table: s.Tables[0],
239+
Unique: true,
240+
Parts: []*schema.IndexPart{
241+
{SeqNo: 0, C: s.Tables[0].Columns[0]},
242+
},
243+
Attrs: []schema.Attr{
244+
&IndexPredicate{P: "id <> 0"},
245+
},
246+
},
247+
}
248+
buf, err := MarshalSpec(s, hclState)
249+
require.NoError(t, err)
250+
const expected = `table "users" {
251+
schema = schema.test
252+
column "id" {
253+
null = false
254+
type = int
255+
auto_increment = true
256+
}
257+
index "index" {
258+
unique = true
259+
columns = [column.id]
260+
where = "id <> 0"
261+
}
262+
}
263+
schema "test" {
264+
}
265+
`
266+
require.EqualValues(t, expected, string(buf))
267+
}
268+
211269
func TestTypes(t *testing.T) {
212270
for _, tt := range []struct {
213271
typeExpr string

0 commit comments

Comments
 (0)