Repository: ariga/atlas
Pull request #726 (2022-04-21 11:07:03)

sql/sqlite: support inspecting generated columns
For more info, see [sqlite/sqlite/src/pragma.c](https://github.com/sqlite/sqlite/blob/171c50ec381c5fb4ff1e8b258a49fbf810917011/src/pragma.c#L1109-L1112).


Patch:
diff --git a/sql/sqlite/driver.go b/sql/sqlite/driver.go
index 50aa76bb..2cd6cd1b 100644
--- a/sql/sqlite/driver.go
+++ b/sql/sqlite/driver.go
@@ -93,3 +93,9 @@ const (
 	TypeText    = "text"    // SQLITE_TYPE_TEXT
 	TypeBlob    = "blob"    // SQLITE_TYPE_BLOB
 )
+
+// SQLite generated columns types.
+const (
+	virtual = "VIRTUAL"
+	stored  = "STORED"
+)
diff --git a/sql/sqlite/inspect.go b/sql/sqlite/inspect.go
index c862b319..69112690 100644
--- a/sql/sqlite/inspect.go
+++ b/sql/sqlite/inspect.go
@@ -127,10 +127,11 @@ func (i *inspect) columns(ctx context.Context, t *schema.Table) error {
 func (i *inspect) addColumn(t *schema.Table, rows *sql.Rows) error {
 	var (
 		nullable, primary   bool
+		hidden              sql.NullInt64
 		name, typ, defaults sql.NullString
 		err                 error
 	)
-	if err = rows.Scan(&name, &typ, &nullable, &defaults, &primary); err != nil {
+	if err = rows.Scan(&name, &typ, &nullable, &defaults, &primary, &hidden); err != nil {
 		return err
 	}
 	c := &schema.Column{
@@ -147,7 +148,13 @@ func (i *inspect) addColumn(t *schema.Table, rows *sql.Rows) error {
 	if defaults.Valid {
 		c.Default = defaultExpr(defaults.String)
 	}
-	// TODO(a8m): extract collation from 'CREATE TABLE' statement.
+	// The hidden flag is set to 2 for VIRTUAL columns, and to
+	// 3 for STORED columns. See: sqlite/pragma.c#sqlite3Pragma.
+	if hidden.Int64 >= 2 {
+		if err := setGenExpr(t, c, hidden.Int64); err != nil {
+			return err
+		}
+	}
 	t.Columns = append(t.Columns, c)
 	if primary {
 		if t.PrimaryKey == nil {
@@ -546,6 +553,33 @@ func autoinc(t *schema.Table) error {
 	return nil
 }
 
+// setGenExpr extracts the generated expression from the CREATE statement
+// and appends it to the column.
+func setGenExpr(t *schema.Table, c *schema.Column, f int64) error {
+	var s CreateStmt
+	if !sqlx.Has(t.Attrs, &s) {
+		return fmt.Errorf("missing CREATE statement for table: %q", t.Name)
+	}
+	re, err := regexp.Compile(fmt.Sprintf("(?i)(?:[(,]\\s*)[\"`]*(%s)[\"`]*[^,]*(?:GENERATED\\s+ALWAYS)*\\s*AS\\s*\\(", c.Name))
+	if err != nil {
+		return err
+	}
+	idx := re.FindAllStringIndex(s.S, 1)
+	if len(idx) != 1 || len(idx[0]) != 2 {
+		return fmt.Errorf("sqlite: generation expression for column %q was not found in create statement", c.Name)
+	}
+	expr := scanExpr(s.S[idx[0][1]-1:])
+	if expr == "" {
+		return fmt.Errorf("sqlite: unexpected empty generation expression for column %q", c.Name)
+	}
+	typ := virtual
+	if f == 3 {
+		typ = stored
+	}
+	c.SetGeneratedExpr(&schema.GeneratedExpr{Expr: expr, Type: typ})
+	return nil
+}
+
 // The following regexes extract named FKs and CHECK constraints defined in table-constraints or inlined
 // as column-constraints. Note, we assume the SQL statements are valid as they are returned by SQLite.
 var (
@@ -631,39 +665,41 @@ func fillChecks(t *schema.Table) error {
 		if len(idx) != 4 {
 			break
 		}
-		var (
-			r, l int
-			expr = c.S[idx[1]-1:]
-		)
-		for i := 0; i < len(expr); i++ {
-			switch expr[i] {
-			case '(':
-				r++
-			case ')':
-				l++
-			case '\'', '"':
-				// Skip unescaped strings.
-				if j := strings.IndexByte(expr[i+1:], expr[i]); j != -1 {
-					i += j + 1
-				}
-			}
-			// Balanced parens.
-			if r == l {
-				check := &schema.Check{Expr: expr[:i+1]}
-				// Matching group for constraint name.
-				if idx[2] != -1 && idx[3] != -1 {
-					check.Name = c.S[idx[2]:idx[3]]
-				}
-				t.Attrs = append(t.Attrs, check)
-				idx[1] = idx[1] + i
-				break
-			}
+		check := &schema.Check{Expr: scanExpr(c.S[idx[1]-1:])}
+		// Matching group for constraint name.
+		if idx[2] != -1 && idx[3] != -1 {
+			check.Name = c.S[idx[2]:idx[3]]
 		}
-		c.S = c.S[idx[1]:]
+		t.Attrs = append(t.Attrs, check)
+		c.S = c.S[idx[1]+len(check.Expr)-1:]
 	}
 	return nil
 }
 
+// scanExpr scans the expression string (wrapped with parens)
+// until its end in the given string. e.g. "(a+1), c int ...".
+func scanExpr(expr string) string {
+	var r, l int
+	for i := 0; i < len(expr); i++ {
+		switch expr[i] {
+		case '(':
+			r++
+		case ')':
+			l++
+		case '\'', '"':
+			// Skip unescaped strings.
+			if j := strings.IndexByte(expr[i+1:], expr[i]); j != -1 {
+				i += j + 1
+			}
+		}
+		// Balanced parens.
+		if r == l {
+			return expr[:i+1]
+		}
+	}
+	return ""
+}
+
 const (
 	// Name of main database file.
 	mainFile = "main"
@@ -672,7 +708,7 @@ const (
 	// Query to list database tables.
 	tablesQuery = "SELECT `name`, `sql` FROM sqlite_master WHERE `type` = 'table' AND `name` NOT LIKE 'sqlite_%'"
 	// Query to list table information.
-	columnsQuery = "SELECT `name`, `type`, (not `notnull`) AS `nullable`, `dflt_value`, (`pk` <> 0) AS `pk`  FROM pragma_table_info('%s') ORDER BY `pk`, `cid`"
+	columnsQuery = "SELECT `name`, `type`, (not `notnull`) AS `nullable`, `dflt_value`, (`pk` <> 0) AS `pk`, `hidden` FROM pragma_table_xinfo('%s') ORDER BY `pk`, `cid`"
 	// Query to list table indexes.
 	indexesQuery = "SELECT `il`.`name`, `il`.`unique`, `il`.`origin`, `il`.`partial`, `m`.`sql` FROM pragma_index_list('%s') AS il JOIN sqlite_master AS m ON il.name = m.name"
 	// Query to list index columns.


Test patch:
diff --git a/sql/sqlite/inspect_test.go b/sql/sqlite/inspect_test.go
index 31fc949d..2c981644 100644
--- a/sql/sqlite/inspect_test.go
+++ b/sql/sqlite/inspect_test.go
@@ -25,22 +25,25 @@ func TestDriver_InspectTable(t *testing.T) {
 		{
 			name: "table columns",
 			before: func(m mock) {
-				m.tableExists("users", true, "CREATE TABLE users(id INTEGER PRIMARY KEY AUTOINCREMENT)")
+				m.tableExists("users", true, "CREATE TABLE users(id INTEGER PRIMARY KEY AUTOINCREMENT, w INT GENERATED ALWAYS AS (a*10), x TEXT AS (typeof(c)) STORED, y TEXT AS (substr(b,a,a+2)))")
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(columnsQuery, "users"))).
 					WillReturnRows(sqltest.Rows(`
- name |   type       | nullable | dflt_value  | primary 
-------+--------------+----------+ ------------+----------
- c1   | int           |  1      |     a       |  0
- c2   | integer       |  0      |     97      |  0
- c3   | varchar(100)  |  1      |    'A'      |  0
- c4   | boolean       |  0      |             |  0
- c5   | json          |  0      |             |  0
- c6   | datetime      |  0      |             |  0
- c7   | blob          |  0      |    x'a'     |  0
- c8   | text          |  0      |             |  0
- c9   | numeric(10,2) |  0      |             |  0
- c10  | real          |  0      |             |  0
- id   | integer       |  0      |     0x1     |  1
+ name |   type       | nullable | dflt_value  | primary  | hidden
+------+--------------+----------+ ------------+----------+----------
+ c1   | int           |  1      |     a       |  0       |  0
+ c2   | integer       |  0      |     97      |  0       |  0
+ c3   | varchar(100)  |  1      |    'A'      |  0       |  0
+ c4   | boolean       |  0      |             |  0       |  0
+ c5   | json          |  0      |             |  0       |  0
+ c6   | datetime      |  0      |             |  0       |  0
+ c7   | blob          |  0      |    x'a'     |  0       |  0
+ c8   | text          |  0      |             |  0       |  0
+ c9   | numeric(10,2) |  0      |             |  0       |  0
+ c10  | real          |  0      |             |  0       |  0
+ w    | int           |  0      |             |  0       |  2
+ x    | text          |  0      |             |  0       |  3
+ y    | text          |  0      |             |  0       |  2
+ id   | integer       |  0      |     0x1     |  1       |  0
 `))
 				m.noIndexes("users")
 				m.noFKs("users")
@@ -58,6 +61,9 @@ 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: "w", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int"}, Raw: "int"}, Attrs: []schema.Attr{&schema.GeneratedExpr{Type: "VIRTUAL", Expr: "(a*10)"}}},
+					{Name: "x", Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}, Raw: "text"}, Attrs: []schema.Attr{&schema.GeneratedExpr{Type: "STORED", Expr: "(typeof(c))"}}},
+					{Name: "y", Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}, Raw: "text"}, Attrs: []schema.Attr{&schema.GeneratedExpr{Type: "VIRTUAL", Expr: "(substr(b,a,a+2))"}}},
 					{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)
@@ -76,10 +82,10 @@ func TestDriver_InspectTable(t *testing.T) {
 				m.tableExists("users", true, "CREATE TABLE users(id INTEGER PRIMARY KEY)")
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(columnsQuery, "users"))).
 					WillReturnRows(sqltest.Rows(`
- name |   type       | nullable | dflt_value  | primary 
-------+--------------+----------+ ------------+----------
- c1   | int           |  1      |             |  0
- c2   | integer       |  0      |             |  0
+ name |   type       | nullable | dflt_value  | primary  | hidden
+------+--------------+----------+ ------------+----------+----------
+ c1   | int           |  1      |             |  0       |  0
+ c2   | integer       |  0      |             |  0       |  0
 `))
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(indexesQuery, "users"))).
 					WillReturnRows(sqltest.Rows(`
@@ -175,11 +181,11 @@ CREATE TABLE users(
 `)
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(columnsQuery, "users"))).
 					WillReturnRows(sqltest.Rows(`
- name |   type       | nullable | dflt_value  | primary 
-------+--------------+----------+ ------------+----------
- c1   | int           |  1      |             |  0
- c2   | integer       |  0      |             |  0
- c3   | integer       |  0      |             |  0
+ name |   type       | nullable | dflt_value  | primary  | hidden
+------+--------------+----------+ ------------+----------+----------
+ c1   | int           |  1      |             |  0       |  0
+ c2   | integer       |  0      |             |  0       |  0
+ c3   | integer       |  0      |             |  0       |  0
 `))
 				m.noIndexes("users")
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(fksQuery, "users"))).