Add migrations to use uast nullable cols in FilePairs

Author: dpordomingo

cli/migrations/README.md

@@ -0,0 +1,17 @@
+### Disclaimer:
+
+These scripts are <u>**not real migrations**</u>. There is no up/down rules. These scripts are not production ready. The idempotence of these scripts is not guaranteed at all..
+
+These scripts were needed to migrate the Database between different states.
+
+- [Add UAST nullable columns](command-UAST-add-columns.go)
+```shell
+# prepares the current "internal.db". Adds "uast_a" and "uast_b" nulable cols
+go run cli/migrations/*.go uast-add-cols internal.db
+```
+
+- [Add UAST to a database, reading from other Database](command-UAST-import-from-source-db.go)
+```shell
+# import UASTs into "internal.db" reading from "source.db"
+go run cli/migrations/*.go uast-import internal.db source.db
+```

cli/migrations/command-UAST-add-columns.go

@@ -0,0 +1,58 @@
+package main
+
+import (
+	"fmt"
+	"log"
+
+	"github.com/src-d/code-annotation/server/dbutil"
+)
+
+type uastColsCmd struct {
+	commandDesc
+	Args struct {
+		InternalDBPath string `description:"filepath to the internal SQLite database"`
+	} `positional-args:"yes" required:"yes"`
+}
+
+var uastColsOpts = uastColsCmd{
+	commandDesc: commandDesc{
+		name:      "uast-add-cols",
+		shortDesc: "Add uast BLOB columns",
+		longDesc:  "Adds 'uast_a' and 'uast_b' BLOB columns to the sqlite://InternalDBPath database",
+	},
+}
+
+const (
+	leftFile  = "a"
+	rightFile = "b"
+
+	addUastQuery = "ALTER TABLE file_pairs ADD COLUMN uast_%s BLOB"
+)
+
+func (c *uastColsCmd) Execute(args []string) error {
+	internalDb, err := dbutil.Open(fmt.Sprintf(sqliteDSN, c.Args.InternalDBPath), false)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	defer internalDb.Close()
+
+	if err := addColumn(internalDb, leftFile); err != nil {
+		return err
+	}
+
+	if err := addColumn(internalDb, rightFile); err != nil {
+		return err
+	}
+
+	fmt.Println("New BLOB columns 'uast_a' and 'uast_b' were added")
+	return nil
+}
+
+func addColumn(db dbutil.DB, side string) error {
+	if _, err := db.Exec(fmt.Sprintf(addUastQuery, side)); err != nil {
+		return err
+	}
+
+	return nil
+}

cli/migrations/command-UAST-import-from-source-db.go

@@ -0,0 +1,176 @@
+package main
+
+import (
+	"fmt"
+	"log"
+
+	"github.com/src-d/code-annotation/server/dbutil"
+)
+
+type uastImportCmd struct {
+	commandDesc
+	Args struct {
+		InternalDBPath string `description:"filepath to the internal SQLite database"`
+		SourceDBPath   string `description:"filepath to the SQLite database containing the UAST to import"`
+	} `positional-args:"yes" required:"yes"`
+}
+
+var uastImportOpts = uastImportCmd{
+	commandDesc: commandDesc{
+		name:      "uast-import",
+		shortDesc: "Import UASTs",
+		longDesc:  "Adds UASTs to the sqlite://InternalDBPath database reading from sqlite://SourceDBPath database",
+	},
+}
+
+const (
+	filePairsWithoutUastQuery = `SELECT
+			blob_id_a, uast_a IS NOT NULL as hasUastA, blob_id_b, uast_b IS NOT NULL as hasUastB
+		FROM file_pairs
+		WHERE hasUastA = 0 or hasUastB = 0`
+	uastByBlobIDQuery = `SELECT uast_%s
+		FROM files
+		WHERE blob_id_%s = CAST($1 AS BLOB)
+		LIMIT 1`
+	updateByBlobIDQuery = `UPDATE file_pairs
+		SET uast_%s = $1
+		WHERE blob_id_%s = $2 and uast_%s IS NULL`
+	indexAddBlobID = `CREATE INDEX blob_id_%s ON file_pairs (blob_id_%s);`
+)
+
+type file struct {
+	side   string
+	blobID string
+}
+
+func (c *uastImportCmd) Execute(args []string) error {
+	internalDb, err := dbutil.Open(fmt.Sprintf(sqliteDSN, c.Args.InternalDBPath), false)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	defer internalDb.Close()
+
+	sourceDb, err := dbutil.Open(fmt.Sprintf(sqliteDSN, c.Args.SourceDBPath), false)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	defer sourceDb.Close()
+
+	fixSourceDb(sourceDb)
+
+	files, fileReadfailures := getFilesToUpdate(internalDb)
+
+	var rowsEdited, uastFails, uastsImported int64
+	for _, file := range files {
+		affectedRows, err := importUastForBlobID(internalDb, sourceDb, file.side, file.blobID)
+		if err != nil {
+			fmt.Println(err)
+			uastFails++
+			continue
+		}
+
+		rowsEdited += affectedRows
+		uastsImported++
+	}
+
+	fmt.Printf(`
+UASTs added into the internal DB:
+- processed files: %d
+  - succeeded:
+    - imported UASTs: %d
+    - edited rows: %d
+  - failed: %d
+- FilePair read errors:  %d
+`, len(files), uastsImported, rowsEdited, uastFails, fileReadfailures)
+
+	if fileReadfailures+uastFails > 0 {
+		log.Fatal("Some rows could not be properly updated")
+	}
+
+	return nil
+}
+
+type files map[string]file
+
+func (f *files) add(blobID string, side string, ignore bool) {
+	if ignore {
+		return
+	}
+
+	if _, ok := (*f)[blobID+"_"+side]; !ok {
+		(*f)[blobID+"_"+side] = file{side: side, blobID: blobID}
+	}
+}
+
+func getFilesToUpdate(internalDb dbutil.DB) (map[string]file, int64) {
+	rows, err := internalDb.Query(filePairsWithoutUastQuery)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	defer rows.Close()
+
+	filesToImport := files{}
+	var failures int64
+	for rows.Next() {
+		var blobIDA, blobIDB string
+		var hasUastA, hasUastB int
+		err := rows.Scan(&blobIDA, &hasUastA, &blobIDB, &hasUastB)
+		if err != nil {
+			fmt.Printf("Failed to read row from internal DB\nerror: %v\n", err)
+			failures++
+			continue
+		}
+
+		filesToImport.add(blobIDA, leftFile, hasUastA == 1)
+		filesToImport.add(blobIDB, rightFile, hasUastB == 1)
+	}
+
+	return filesToImport, failures
+}
+
+func importUastForBlobID(internalDb dbutil.DB, sourceDb dbutil.DB, side string, blobID string) (int64, error) {
+	uast, err := getUastByBlobID(sourceDb, side, blobID)
+	if err != nil {
+		return 0, fmt.Errorf("uast_%s could not be retrieved for blobID#%s; %s", side, blobID, err)
+	}
+
+	return setUastToBlobID(internalDb, side, blobID, uast)
+}
+
+func getUastByBlobID(sourceDb dbutil.DB, side string, blobID string) ([]byte, error) {
+	var uast []byte
+	if err := sourceDb.QueryRow(fmt.Sprintf(uastByBlobIDQuery, side, side), blobID).Scan(&uast); err != nil {
+		return nil, err
+	}
+
+	return uast, nil
+}
+
+func setUastToBlobID(internalDb dbutil.DB, side string, blobID string, uast []byte) (int64, error) {
+	res, err := internalDb.Exec(fmt.Sprintf(updateByBlobIDQuery, side, side, side), uast, blobID)
+	if err != nil {
+		return 0, fmt.Errorf("uast_%s could not be saved for blobID#%s; %s", side, blobID, err)
+	}
+
+	rows, _ := res.RowsAffected()
+	if rows == 0 {
+		return 0, fmt.Errorf("no uast_%s to be imported for blobID#%s", side, blobID)
+	}
+
+	return rows, nil
+}
+
+func fixSourceDb(sourceDb dbutil.DB) error {
+	if _, err := sourceDb.Exec(fmt.Sprintf(indexAddBlobID, leftFile, leftFile)); err != nil {
+		return fmt.Errorf("can not create index over blob_id_%s; %s", leftFile, err)
+	}
+
+	if _, err := sourceDb.Exec(fmt.Sprintf(indexAddBlobID, rightFile, rightFile)); err != nil {
+		return fmt.Errorf("can not create index over blob_id_%s; %s", rightFile, err)
+	}
+
+	return nil
+}

cli/migrations/command.go

@@ -0,0 +1,31 @@
+package main
+
+import (
+	_ "net/http/pprof"
+)
+
+// Command is a runnable command
+type Command interface {
+	Name() string
+	ShortDesc() string
+	LongDesc() string
+	Execute(args []string) error
+}
+
+type commandDesc struct {
+	name      string
+	shortDesc string
+	longDesc  string
+}
+
+func (c *commandDesc) Name() string {
+	return c.name
+}
+
+func (c *commandDesc) ShortDesc() string {
+	return c.shortDesc
+}
+
+func (c *commandDesc) LongDesc() string {
+	return c.longDesc
+}

cli/migrations/main.go

@@ -0,0 +1,43 @@
+package main
+
+import (
+	"fmt"
+	_ "net/http/pprof"
+	"os"
+
+	"github.com/jessevdk/go-flags"
+)
+
+const (
+	description string = "Migrate internal database"
+	sqliteDSN          = "sqlite://%s"
+)
+
+func main() {
+	parser := flags.NewParser(nil, flags.Default)
+	addCommand(parser, &uastColsOpts)
+	addCommand(parser, &uastImportOpts)
+	parse(parser, description)
+}
+
+func addCommand(parser *flags.Parser, command Command) {
+	if _, err := parser.AddCommand(command.Name(), command.ShortDesc(), command.LongDesc(), command); err != nil {
+		panic(err)
+	}
+}
+
+func parse(parser *flags.Parser, description string) {
+	parser.LongDescription = description
+	if _, err := parser.Parse(); err != nil {
+		if err, ok := err.(*flags.Error); ok {
+			if err.Type == flags.ErrHelp {
+				os.Exit(0)
+			}
+
+			fmt.Println()
+			parser.WriteHelp(os.Stdout)
+		}
+
+		os.Exit(1)
+	}
+}