Repository: ariga/atlas
Pull request #244 (2021-11-24 10:12:00)

sql/sqlite: detect AUTOINCREMENT clause in primary-keys


Patch:
diff --git a/sql/sqlite/driver.go b/sql/sqlite/driver.go
index 3f17ed02..ee108ace 100644
--- a/sql/sqlite/driver.go
+++ b/sql/sqlite/driver.go
@@ -8,6 +8,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"regexp"
 	"strconv"
 	"strings"
 	"unicode"
@@ -151,6 +152,7 @@ func (d *Driver) columns(ctx context.Context, t *schema.Table) error {
 			return fmt.Errorf("sqlite: %w", err)
 		}
 	}
+	autoinc(t)
 	return rows.Err()
 }
 
@@ -491,6 +493,12 @@ type (
 		S string
 	}
 
+	// AutoIncrement describes the `AUTOINCREMENT` configuration.
+	// https://www.sqlite.org/autoinc.html
+	AutoIncrement struct {
+		schema.Attr
+	}
+
 	// WithoutRowID describes the `WITHOUT ROWID` configuration.
 	// See: https://sqlite.org/withoutrowid.html
 	WithoutRowID struct {
@@ -541,6 +549,23 @@ func isNumber(s string) bool {
 	return true
 }
 
+var reAutoinc = regexp.MustCompile(`(?i)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) {
+	if t.PrimaryKey == nil || len(t.PrimaryKey.Parts) != 1 || t.PrimaryKey.Parts[0].C == nil {
+		return
+	}
+	if c := (CreateStmt{}); !sqlx.Has(t.Attrs, &c) || !reAutoinc.MatchString(c.S) {
+		return
+	}
+	// 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{})
+}
+
 const (
 	// Query to list attached database files.
 	databasesQuery = "SELECT `name`, `file` FROM pragma_database_list()"
diff --git a/sql/sqlite/migrate.go b/sql/sqlite/migrate.go
index 550cb4f9..c902be32 100644
--- a/sql/sqlite/migrate.go
+++ b/sql/sqlite/migrate.go
@@ -49,7 +49,8 @@ func (m *migrate) addTable(ctx context.Context, add *schema.AddTable) error {
 		b.MapComma(add.T.Columns, func(i int, b *sqlx.Builder) {
 			m.column(b, add.T.Columns[i])
 		})
-		if pk := add.T.PrimaryKey; pk != nil {
+		// Primary keys with auto-increment are inlined on the column definition.
+		if pk := add.T.PrimaryKey; pk != nil && !autoincPK(pk) {
 			b.Comma().P("PRIMARY KEY")
 			m.indexParts(b, pk.Parts)
 		}
@@ -123,6 +124,9 @@ func (m *migrate) column(b *sqlx.Builder, c *schema.Column) {
 	if x, ok := c.Default.(*schema.RawExpr); ok {
 		b.P("DEFAULT", x.X)
 	}
+	if sqlx.Has(c.Attrs, &AutoIncrement{}) {
+		b.P("PRIMARY KEY AUTOINCREMENT")
+	}
 }
 
 func (m *migrate) addIndexes(ctx context.Context, t *schema.Table, indexes ...*schema.Index) error {
@@ -328,6 +332,11 @@ func alterable(modify *schema.ModifyTable) bool {
 	return true
 }
 
+func autoincPK(pk *schema.Index) bool {
+	return sqlx.Has(pk.Attrs, &AutoIncrement{}) ||
+		len(pk.Parts) == 1 && pk.Parts[0].C != nil && sqlx.Has(pk.Parts[0].C.Attrs, &AutoIncrement{})
+}
+
 // Build instantiates a new builder and writes the given phrase to it.
 func Build(phrase string) *sqlx.Builder {
 	b := &sqlx.Builder{QuoteChar: '`'}


Test patch:
diff --git a/internal/integration/sqlite_test.go b/internal/integration/sqlite_test.go
index 0b1f43c5..358fdcf1 100644
--- a/internal/integration/sqlite_test.go
+++ b/internal/integration/sqlite_test.go
@@ -134,6 +134,26 @@ func TestSQLite_AddIndexedColumns(t *testing.T) {
 	})
 }
 
+func TestSQLite_AutoIncrement(t *testing.T) {
+	liteRun(t, func(t *liteTest) {
+		usersT := &schema.Table{
+			Name: "users",
+			Columns: []*schema.Column{
+				{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}}, Attrs: []schema.Attr{sqlite.AutoIncrement{}}},
+			},
+		}
+		usersT.PrimaryKey = &schema.Index{Parts: []*schema.IndexPart{{C: usersT.Columns[0]}}}
+		t.migrate(&schema.AddTable{T: usersT})
+		t.dropTables(usersT.Name)
+		_, err := t.db.Exec("INSERT INTO users DEFAULT VALUES")
+		require.NoError(t, err)
+		var id int
+		err = t.db.QueryRow("SELECT id FROM users").Scan(&id)
+		require.NoError(t, err)
+		require.Equal(t, 1, id)
+	})
+}
+
 func (t *liteTest) loadRealm() *schema.Realm {
 	r, err := t.drv.InspectRealm(context.Background(), &schema.InspectRealmOption{
 		Schemas: []string{"main"},
@@ -164,9 +184,8 @@ func (t *liteTest) users() *schema.Table {
 		Schema: t.realm().Schemas[0],
 		Columns: []*schema.Column{
 			{
-				Name:  "id",
-				Type:  &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}},
-				Attrs: []schema.Attr{&postgres.Identity{}},
+				Name: "id",
+				Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}},
 			},
 			{
 				Name: "x",
diff --git a/sql/sqlite/driver_test.go b/sql/sqlite/driver_test.go
index 0a2add05..8cd24f79 100644
--- a/sql/sqlite/driver_test.go
+++ b/sql/sqlite/driver_test.go
@@ -39,7 +39,7 @@ func TestDriver_InspectTable(t *testing.T) {
 			name: "table columns",
 			before: func(m mock) {
 				m.systemVars("3.36.0")
-				m.tableExists("users", true, "CREATE TABLE users(id INTEGER PRIMARY KEY)")
+				m.tableExists("users", true, "CREATE TABLE users(id INTEGER PRIMARY KEY AUTOINCREMENT)")
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(columnsQuery, "users"))).
 					WillReturnRows(sqltest.Rows(`
  name |   type       | nullable | dflt_value  | primary 
@@ -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"}},
+					{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}, Raw: "integer"}, Attrs: []schema.Attr{AutoIncrement{}}},
 				}
 				require.Equal(t.Columns, columns)
 				require.EqualValues(&schema.Index{
@@ -80,6 +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{}},
 				}, t.PrimaryKey)
 			},
 		},