Repository: ariga/atlas
Pull request #1447 (2023-02-05 13:48:03)

sql/sqlite: support reading index part with expressions
Fixed https://github.com/ariga/atlas/issues/1436


Patch:
diff --git a/sql/sqlite/inspect.go b/sql/sqlite/inspect.go
index e6de5953..2ebb07b2 100644
--- a/sql/sqlite/inspect.go
+++ b/sql/sqlite/inspect.go
@@ -227,8 +227,11 @@ func (i *inspect) addIndexes(t *schema.Table, rows *sql.Rows) error {
 	return nil
 }
 
-// A regexp to extract index parts.
-var reIdxParts = regexp.MustCompile("(?i)ON\\s+[\"`]*(?:\\w+)[\"`]*\\s*\\((.+)\\)")
+var (
+	// A regexp to extract index parts.
+	reIdxParts = regexp.MustCompile("(?i)ON\\s+[\"`]*(?:\\w+)[\"`]*\\s*\\((.+?)\\)(\\s*WHERE\\s+.+)?$")
+	reIdxDesc  = regexp.MustCompile("(?i)\\s+DESC\\s*$")
+)
 
 func (i *inspect) indexInfo(ctx context.Context, t *schema.Table, idx *schema.Index) error {
 	var (
@@ -271,15 +274,23 @@ func (i *inspect) indexInfo(ctx context.Context, t *schema.Table, idx *schema.In
 	if !sqlx.Has(idx.Attrs, &c) || !reIdxParts.MatchString(c.S) {
 		return nil
 	}
-	parts := strings.Split(reIdxParts.FindStringSubmatch(c.S)[1], ",")
-	// Unable to parse index parts correctly.
-	if len(parts) != len(idx.Parts) {
-		return nil
-	}
-	for i, p := range idx.Parts {
+	x := reIdxParts.FindStringSubmatch(c.S)[1]
+	for _, p := range idx.Parts {
+		j := sqlx.ExprLastIndex(x)
+		// Unable to parse index parts correctly.
+		if j == -1 {
+			return nil
+		}
 		if p.X != nil {
-			p.X.(*schema.RawExpr).X = strings.TrimSpace(parts[i])
+			// Remove any extra spaces and the "DESC" clause
+			// in case the key-part is descending.
+			kx := strings.TrimSpace(x[:j+1])
+			if p.Desc {
+				kx = reIdxDesc.ReplaceAllString(kx, "")
+			}
+			p.X.(*schema.RawExpr).X = kx
 		}
+		x = strings.TrimLeft(x[j+1:], ", ")
 	}
 	return nil
 }


Test patch:
diff --git a/internal/integration/testdata/sqlite/index-expr.txt b/internal/integration/testdata/sqlite/index-expr.txt
index d2d9337d..d31fe5d8 100644
--- a/internal/integration/testdata/sqlite/index-expr.txt
+++ b/internal/integration/testdata/sqlite/index-expr.txt
@@ -39,11 +39,11 @@ table "users" {
   }
   index "full_name" {
     on {
-        expr = "first_name || '''s first name'"
+        expr = "lower(first_name) || '''s first name'"
     }
   }
 }
 
 -- 2.sql --
 CREATE TABLE "users" (`first_name` text NOT NULL)
-CREATE INDEX `full_name` ON `users` ((first_name || '''s first name'))
\ No newline at end of file
+CREATE INDEX `full_name` ON `users` ((lower(first_name) || '''s first name'))
\ No newline at end of file
diff --git a/sql/internal/sqlx/sqlx_test.go b/sql/internal/sqlx/sqlx_test.go
index 72251621..96708802 100644
--- a/sql/internal/sqlx/sqlx_test.go
+++ b/sql/internal/sqlx/sqlx_test.go
@@ -112,7 +112,6 @@ func TestMayWrap(t *testing.T) {
 				expect = "(" + expect + ")"
 			}
 			require.Equal(t, expect, MayWrap(tt.input))
-
 		})
 	}
 }
diff --git a/sql/sqlite/inspect_test.go b/sql/sqlite/inspect_test.go
index 81477a21..b5b0be7d 100644
--- a/sql/sqlite/inspect_test.go
+++ b/sql/sqlite/inspect_test.go
@@ -86,6 +86,7 @@ func TestDriver_InspectTable(t *testing.T) {
 ------+--------------+----------+ ------------+----------+----------
  c1   | int           |  1      |             |  0       |  0
  c2   | integer       |  0      |             |  0       |  0
+ c3   | json          |  0      |             |  0       |  0
 `))
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(indexesQuery, "users"))).
 					WillReturnRows(sqltest.Rows(`
@@ -94,6 +95,7 @@ func TestDriver_InspectTable(t *testing.T) {
  c1u   |  1           |  c     |  0       | CREATE UNIQUE INDEX c1u on users(c1, c2)
  c1_c2 |  0           |  c     |  1       | CREATE INDEX c1_c2 on users(c1, c2*2) WHERE c1 <> NULL
  c1_x  |  0           |  c     |  0       | CREATE INDEX c1_x ON users (f(c1))
+ c3_x  |  0           |  c     |  0       | CREATE INDEX c3_x ON users (json_extract(c3, '$.x') desc, json_extract(c3, '$.y') desc)
 `))
 				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(indexColumnsQuery, "c1u"))).
 					WillReturnRows(sqltest.Rows(`
@@ -114,6 +116,13 @@ func TestDriver_InspectTable(t *testing.T) {
  name  |   desc |
 -------+--------+
  nil   |  0     |
+`))
+				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(indexColumnsQuery, "c3_x"))).
+					WillReturnRows(sqltest.Rows(`
+ name  |   desc |
+-------+--------+
+ nil   |  1     |
+ nil   |  1     |
 `))
 				m.noFKs("users")
 			},
@@ -122,6 +131,7 @@ func TestDriver_InspectTable(t *testing.T) {
 				columns := []*schema.Column{
 					{Name: "c1", Type: &schema.ColumnType{Null: true, Type: &schema.IntegerType{T: "int"}, Raw: "int"}},
 					{Name: "c2", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "integer"}, Raw: "integer"}},
+					{Name: "c3", Type: &schema.ColumnType{Type: &schema.JSONType{T: "json"}, Raw: "json"}},
 				}
 				indexes := []*schema.Index{
 					{
@@ -161,6 +171,18 @@ func TestDriver_InspectTable(t *testing.T) {
 							&IndexOrigin{O: "c"},
 						},
 					},
+					{
+						Name:  "c3_x",
+						Table: t,
+						Parts: []*schema.IndexPart{
+							{SeqNo: 1, X: &schema.RawExpr{X: "json_extract(c3, '$.x')"}, Desc: true},
+							{SeqNo: 2, X: &schema.RawExpr{X: "json_extract(c3, '$.y')"}, Desc: true},
+						},
+						Attrs: []schema.Attr{
+							&CreateStmt{S: "CREATE INDEX c3_x ON users (json_extract(c3, '$.x') desc, json_extract(c3, '$.y') desc)"},
+							&IndexOrigin{O: "c"},
+						},
+					},
 				}
 				require.Equal(t.Columns, columns)
 				require.Equal(t.Indexes, indexes)