Repository: ariga/atlas
Pull request #363 (2022-01-03 14:12:21)

sql/sqlite: support creating tables with custom sequence


Patch:
diff --git a/sql/sqlite/inspect.go b/sql/sqlite/inspect.go
index 5f4cc6dd..540fced8 100644
--- a/sql/sqlite/inspect.go
+++ b/sql/sqlite/inspect.go
@@ -129,8 +129,7 @@ func (i *inspect) columns(ctx context.Context, t *schema.Table) error {
 			return fmt.Errorf("sqlite: %w", err)
 		}
 	}
-	autoinc(t)
-	return rows.Err()
+	return autoinc(t)
 }
 
 // addColumn scans the current row and adds a new column from it to the table.
@@ -471,6 +470,13 @@ type (
 	// https://www.sqlite.org/autoinc.html
 	AutoIncrement struct {
 		schema.Attr
+		// Seq represents the value in sqlite_sequence table.
+		// i.e. https://www.sqlite.org/fileformat2.html#seqtab.
+		//
+		// Setting this value manually to > 0 indicates that
+		// a custom value is necessary and should be handled
+		// on migrate.
+		Seq int64
 	}
 
 	// WithoutRowID describes the `WITHOUT ROWID` configuration.
@@ -547,21 +553,35 @@ func isBlob(s string) bool {
 	return false
 }
 
-var reAutoinc = regexp.MustCompile(`(?i)PRIMARY\s+KEY\s+AUTOINCREMENT`)
+var reAutoinc = regexp.MustCompile("(?i)(?:[(,]\\s*)[\"`]?(\\w+)[\"`]?\\s+INTEGER\\s+[^,]*PRIMARY\\s+KEY\\s+[^,]*AUTOINCREMENT")
 
 // autoinc checks if the table contains a "PRIMARY KEY AUTOINCREMENT" on its
 // CREATE statement, according to https://www.sqlite.org/syntax/column-constraint.html.
 // This is a workaround until we will embed a proper SQLite parser in atlas.
-func autoinc(t *schema.Table) {
+func autoinc(t *schema.Table) error {
+	var c CreateStmt
+	if !sqlx.Has(t.Attrs, &c) {
+		return fmt.Errorf("missing CREATE statment for table: %q", t.Name)
+	}
 	if t.PrimaryKey == nil || len(t.PrimaryKey.Parts) != 1 || t.PrimaryKey.Parts[0].C == nil {
-		return
+		return nil
 	}
-	if c := (CreateStmt{}); !sqlx.Has(t.Attrs, &c) || !reAutoinc.MatchString(c.S) {
-		return
+	matches := reAutoinc.FindStringSubmatch(c.S)
+	if len(matches) != 2 {
+		return nil
 	}
+	pkc, ok := t.Column(matches[1])
+	if !ok {
+		return fmt.Errorf("sqlite: column %q was not found for AUTOINCREMENT", matches[1])
+	}
+	if t.PrimaryKey == nil || len(t.PrimaryKey.Parts) != 1 || t.PrimaryKey.Parts[0].C != pkc {
+		return fmt.Errorf("sqlite: unexpected primary key: %v", t.PrimaryKey)
+	}
+	inc := &AutoIncrement{}
 	// Annotate table elements with "AUTOINCREMENT".
-	t.PrimaryKey.Attrs = append(t.PrimaryKey.Attrs, AutoIncrement{})
-	t.PrimaryKey.Parts[0].C.Attrs = append(t.PrimaryKey.Parts[0].C.Attrs, AutoIncrement{})
+	t.PrimaryKey.Attrs = append(t.PrimaryKey.Attrs, inc)
+	pkc.Attrs = append(pkc.Attrs, inc)
+	return nil
 }
 
 // The following regexes extract named FKs and CHECK constraints defined in table-constraints or inlined
diff --git a/sql/sqlite/migrate.go b/sql/sqlite/migrate.go
index f9341b81..35b02028 100644
--- a/sql/sqlite/migrate.go
+++ b/sql/sqlite/migrate.go
@@ -6,6 +6,8 @@ package sqlite
 
 import (
 	"context"
+	"database/sql"
+	"errors"
 	"fmt"
 	"strings"
 
@@ -27,7 +29,7 @@ func (p *planApply) PlanChanges(ctx context.Context, name string, changes []sche
 			Transactional: true,
 		},
 	}
-	if err := s.plan(changes); err != nil {
+	if err := s.plan(ctx, changes); err != nil {
 		return nil, err
 	}
 	for _, c := range s.Changes {
@@ -56,15 +58,15 @@ type state struct {
 
 // Exec executes the changes on the database. An error is returned
 // if one of the operations fail, or a change is not supported.
-func (s *state) plan(changes []schema.Change) (err error) {
+func (s *state) plan(ctx context.Context, changes []schema.Change) (err error) {
 	for _, c := range changes {
 		switch c := c.(type) {
 		case *schema.AddTable:
-			err = s.addTable(c)
+			err = s.addTable(ctx, c)
 		case *schema.DropTable:
 			err = s.dropTable(c)
 		case *schema.ModifyTable:
-			err = s.modifyTable(c)
+			err = s.modifyTable(ctx, c)
 		default:
 			err = fmt.Errorf("unsupported change %T", c)
 		}
@@ -84,7 +86,7 @@ func (s *state) plan(changes []schema.Change) (err error) {
 }
 
 // addTable builds and executes the query for creating a table in a schema.
-func (s *state) addTable(add *schema.AddTable) error {
+func (s *state) addTable(ctx context.Context, add *schema.AddTable) error {
 	b := Build("CREATE TABLE").Ident(add.T.Name)
 	if sqlx.Has(add.Extra, &schema.IfNotExists{}) {
 		b.P("IF NOT EXISTS")
@@ -118,6 +120,9 @@ func (s *state) addTable(add *schema.AddTable) error {
 		Reverse: Build("DROP TABLE").Table(add.T).String(),
 		Comment: fmt.Sprintf("create %q table", add.T.Name),
 	})
+	if err := s.tableSeq(ctx, add); err != nil {
+		return err
+	}
 	return s.addIndexes(add.T, add.T.Indexes...)
 }
 
@@ -140,7 +145,7 @@ func (s *state) dropTable(drop *schema.DropTable) error {
 // If the modification contains changes that are not index creation/deletion or a simple column
 // addition, the changes are applied using a temporary table following the procedure mentioned
 // in: https://www.sqlite.org/lang_altertable.html#making_other_kinds_of_table_schema_changes.
-func (s *state) modifyTable(modify *schema.ModifyTable) error {
+func (s *state) modifyTable(ctx context.Context, modify *schema.ModifyTable) error {
 	if alterable(modify) {
 		return s.alterTable(modify)
 	}
@@ -150,7 +155,7 @@ func (s *state) modifyTable(modify *schema.ModifyTable) error {
 	newT.Indexes = nil
 	newT.Name = "new_" + newT.Name
 	// Create a new table with a temporary name, and copy the existing rows to it.
-	if err := s.addTable(&schema.AddTable{T: &newT}); err != nil {
+	if err := s.addTable(ctx, &schema.AddTable{T: &newT}); err != nil {
 		return err
 	}
 	if err := s.copyRows(modify.T, &newT, modify.Changes); err != nil {
@@ -361,6 +366,37 @@ func (s *state) alterTable(modify *schema.ModifyTable) error {
 	return nil
 }
 
+// tableSeq sets the sequence value of the table if it was provided by
+// the user on table creation.
+func (s *state) tableSeq(ctx context.Context, add *schema.AddTable) error {
+	var inc AutoIncrement
+	for _, c := range add.T.Columns {
+		if sqlx.Has(c.Attrs, &inc) {
+			break
+		}
+	}
+	if inc.Seq == 0 {
+		return nil
+	}
+	// SQLite tracks the AUTOINCREMENT in the "sqlite_sequence" table that is created and initialized automatically
+	// whenever such table is created. However, the rows are populated after the first insertion, and therefore, we
+	// check if the sequence row exists, and we either created it for non zero sequence, or update it if it is 0 and
+	// the configured sequence is > 0.
+	var seq int64
+	switch err := s.QueryRowContext(ctx, "SELECT seq FROM sqlite_sequence WHERE name = ?", add.T.Name).Scan(&seq); {
+	case errors.Is(err, sql.ErrNoRows):
+		s.append(&migrate.Change{
+			Cmd:     fmt.Sprintf("INSERT INTO sqlite_sequence (name, seq) VALUES (%q, %d)", add.T.Name, inc.Seq),
+			Source:  add,
+			Reverse: fmt.Sprintf("UPDATE sqlite_sequence SET seq = 0 WHERE name = %q", add.T.Name),
+			Comment: fmt.Sprintf("set sequence for %q table", add.T.Name),
+		})
+	case err != nil:
+		return fmt.Errorf("query sequence value for table %q: %w", add.T.Name, err)
+	}
+	return nil
+}
+
 func (s *state) append(c *migrate.Change) {
 	s.Changes = append(s.Changes, c)
 }


Test patch:
diff --git a/sql/sqlite/inspect_test.go b/sql/sqlite/inspect_test.go
index 027d6293..11fda179 100644
--- a/sql/sqlite/inspect_test.go
+++ b/sql/sqlite/inspect_test.go
@@ -72,7 +72,7 @@ func TestDriver_InspectTable(t *testing.T) {
 					{Name: "c8", Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}, Raw: "text"}},
 					{Name: "c9", Type: &schema.ColumnType{Type: &schema.DecimalType{T: "numeric", Precision: 10, Scale: 2}, Raw: "numeric(10,2)"}},
 					{Name: "c10", Type: &schema.ColumnType{Type: &schema.FloatType{T: "real"}, Raw: "real"}},
-					{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}, Raw: "integer"}, Attrs: []schema.Attr{AutoIncrement{}}, Default: &schema.Literal{V: "0x1"}},
+					{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}, Raw: "integer"}, Attrs: []schema.Attr{&AutoIncrement{}}, Default: &schema.Literal{V: "0x1"}},
 				}
 				require.Equal(t.Columns, columns)
 				require.EqualValues(&schema.Index{
@@ -80,7 +80,7 @@ func TestDriver_InspectTable(t *testing.T) {
 					Unique: true,
 					Table:  t,
 					Parts:  []*schema.IndexPart{{SeqNo: 1, C: columns[len(columns)-1]}},
-					Attrs:  []schema.Attr{AutoIncrement{}},
+					Attrs:  []schema.Attr{&AutoIncrement{}},
 				}, t.PrimaryKey)
 			},
 		},
@@ -446,6 +446,9 @@ func (m mock) systemVars(version string) {
       NOCASE
       BINARY
 `))
+}
+
+func (m mock) tableExists(table string, exists bool, stmt ...string) {
 	m.ExpectQuery(sqltest.Escape(databasesQuery + " WHERE name IN (?)")).
 		WithArgs("main").
 		WillReturnRows(sqltest.Rows(`
@@ -453,9 +456,6 @@ func (m mock) systemVars(version string) {
 ------+-----------
  main |   
 `))
-}
-
-func (m mock) tableExists(table string, exists bool, stmt ...string) {
 	rows := sqlmock.NewRows([]string{"name", "sql"})
 	if exists {
 		rows.AddRow(table, stmt[0])
diff --git a/sql/sqlite/migrate_test.go b/sql/sqlite/migrate_test.go
index 1e140941..61b20ba7 100644
--- a/sql/sqlite/migrate_test.go
+++ b/sql/sqlite/migrate_test.go
@@ -6,8 +6,11 @@ package sqlite
 
 import (
 	"context"
+	"database/sql"
 	"testing"
 
+	"ariga.io/atlas/sql/internal/sqltest"
+
 	"ariga.io/atlas/sql/migrate"
 	"ariga.io/atlas/sql/schema"
 	"github.com/DATA-DOG/go-sqlmock"
@@ -17,6 +20,7 @@ import (
 func TestPlanChanges(t *testing.T) {
 	tests := []struct {
 		changes []schema.Change
+		mock    func(mock)
 		plan    *migrate.Plan
 	}{
 		{
@@ -41,6 +45,36 @@ func TestPlanChanges(t *testing.T) {
 				Changes:       []*migrate.Change{{Cmd: "CREATE TABLE `posts` (`id` integer NOT NULL, `text` text NULL, CHECK (text <> ''), CONSTRAINT `positive_id` CHECK (id <> 0))", Reverse: "DROP TABLE `posts`"}},
 			},
 		},
+		{
+			changes: []schema.Change{
+				&schema.AddTable{
+					T: &schema.Table{
+						Name: "posts",
+						Columns: []*schema.Column{
+							{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}}, Attrs: []schema.Attr{&AutoIncrement{Seq: 1024}}},
+							{Name: "text", Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}, Null: true}},
+						},
+						Attrs: []schema.Attr{
+							&schema.Check{Expr: "(text <> '')"},
+							&schema.Check{Name: "positive_id", Expr: "(id <> 0)"},
+						},
+					},
+				},
+			},
+			mock: func(m mock) {
+				m.ExpectQuery(sqltest.Escape("SELECT seq FROM sqlite_sequence WHERE name = ?")).
+					WithArgs("posts").
+					WillReturnError(sql.ErrNoRows)
+			},
+			plan: &migrate.Plan{
+				Reversible:    true,
+				Transactional: true,
+				Changes: []*migrate.Change{
+					{Cmd: "CREATE TABLE `posts` (`id` integer NOT NULL PRIMARY KEY AUTOINCREMENT, `text` text NULL, CHECK (text <> ''), CONSTRAINT `positive_id` CHECK (id <> 0))", Reverse: "DROP TABLE `posts`"},
+					{Cmd: `INSERT INTO sqlite_sequence (name, seq) VALUES ("posts", 1024)`, Reverse: `UPDATE sqlite_sequence SET seq = 0 WHERE name = "posts"`},
+				},
+			},
+		},
 		{
 			changes: []schema.Change{
 				&schema.DropTable{T: &schema.Table{Name: "posts"}},
@@ -135,6 +169,9 @@ func TestPlanChanges(t *testing.T) {
 		require.NoError(t, err)
 		m := mock{mk}
 		m.systemVars("3.36.0")
+		if tt.mock != nil {
+			tt.mock(m)
+		}
 		drv, err := Open(db)
 		require.NoError(t, err)
 		plan, err := drv.PlanChanges(context.Background(), "plan", tt.changes)