yvbolt/db_sqlite.go

package main

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"path/filepath"
	"strings"
	"unsafe"

	_ "yvbolt/sqliteclidriver"

	"github.com/ying32/govcl/vcl"
)

const (
	sqliteTablesCaption = "Tables"
	sqliteFilter        = "SQLite database|*.db;*.db3;*.sqlite;*.sqlite3|All files|*.*"
)

type sqliteLoadedDatabase struct {
	displayName string
	path        string
	db          *sql.DB
	nav         *vcl.TTreeNode

	arena []*navData // keepalive
}

func (ld *sqliteLoadedDatabase) DisplayName() string {
	return ld.displayName
}

func (ld *sqliteLoadedDatabase) RootElement() *vcl.TTreeNode {
	return ld.nav
}

func (ld *sqliteLoadedDatabase) Keepalive(ndata *navData) {
	ld.arena = append(ld.arena, ndata)
}

func (ld *sqliteLoadedDatabase) RenderForNav(f *TMainForm, ndata *navData) error {

	if len(ndata.bucketPath) == 0 {
		// Top-level
		f.propertiesBox.SetText("Please select...")
		return nil

	} else if len(ndata.bucketPath) == 1 {
		// Category (tables, ...)
		f.propertiesBox.SetText("Please select...")
		return nil

	} else if len(ndata.bucketPath) == 2 && ndata.bucketPath[0] == sqliteTablesCaption {
		// Render for specific table
		tableName := ndata.bucketPath[1]

		// Get some basic properties
		r := ld.db.QueryRow(`SELECT sql FROM sqlite_schema WHERE name = ?;`, tableName)
		var schemaStmt string
		err := r.Scan(&schemaStmt)
		if err != nil {
			schemaStmt = fmt.Sprintf("* Failed to describe table %q: %s", tableName, err.Error())
		}

		// Display table properties
		f.propertiesBox.SetText(fmt.Sprintf("Selected table %q\n\nSchema:\n\n%s", tableName, schemaStmt))

		// Load column details
		// Use SELECT form instead of common PRAGMA table_info so we can just get names
		// We could possibly get this from the main data select, but this will
		// work even when there are 0 results
		columnNames, err := ld.sqliteGetColumnNamesForTable(tableName)
		if err != nil {
			return fmt.Errorf("Failed to load columns for table %q: %w", tableName)
		}

		populateColumns(columnNames, f.contentBox)

		// Select count(*) so we know to display a warning if there are too many entries
		// TODO

		// Select * with small limit
		datar, err := ld.db.Query(`SELECT * FROM "` + tableName + `" LIMIT 1000`) // WARNING can't prepare this parameter, but it comes from the DB (trusted)
		if err != nil {
			return fmt.Errorf("Failed to load data for table %q: %w", tableName, err)
		}
		defer datar.Close()
		populateRows(datar, f.contentBox)

		// We successfully populated the data grid
		vcl_stringgrid_columnwidths(f.contentBox)
		f.contentBox.SetEnabled(true)
		return nil

	} else {
		// ??? unknown
		return errors.New("?")

	}

}

func (ld *sqliteLoadedDatabase) sqliteGetColumnNamesForTable(tableName string) ([]string, error) {
	colr, err := ld.db.Query(`SELECT name FROM pragma_table_info( ? )`, tableName)
	if err != nil {
		return nil, fmt.Errorf("Query: %w", err)
	}
	defer colr.Close()

	var ret []string

	for colr.Next() {

		var columnName string
		err = colr.Scan(&columnName)
		if err != nil {
			return nil, fmt.Errorf("Scan: %w", colr.Err())
		}

		ret = append(ret, columnName)
	}
	if colr.Err() != nil {
		return nil, colr.Err()
	}

	return ret, nil
}

func populateColumns(names []string, dest *vcl.TStringGrid) {
	for _, columnName := range names {
		col := dest.Columns().Add()
		col.Title().SetCaption(columnName)
	}
}

func populateRows(rr *sql.Rows, dest *vcl.TStringGrid) {

	numColumns := int(dest.Columns().Count())

	for rr.Next() {
		fields := make([]interface{}, numColumns)
		pfields := make([]interface{}, numColumns)
		for i := 0; i < numColumns; i += 1 {
			pfields[i] = &fields[i]
		}

		err := rr.Scan(pfields...)
		if err != nil {
			vcl.ShowMessageFmt("Failed to load data: %s", err.Error())
			return
		}

		rpos := dest.RowCount()
		dest.SetRowCount(rpos + 1)
		for i := 0; i < len(fields); i += 1 {
			dest.SetCells(int32(i), rpos, formatAny(fields[i]))
		}

	}
	if rr.Err() != nil {
		vcl.ShowMessageFmt("Failed to load data: %s", rr.Err().Error())
		return
	}
}

func (ld *sqliteLoadedDatabase) ExecQuery(query string, resultArea *vcl.TStringGrid) error {
	rr, err := ld.db.Query(query)
	if err != nil {
		return err
	}

	defer rr.Close()

	vcl_stringgrid_clear(resultArea)

	columns, err := rr.Columns()
	if err != nil {
		return err
	}

	populateColumns(columns, resultArea)

	populateRows(rr, resultArea)

	vcl_stringgrid_columnwidths(resultArea)
	resultArea.SetEnabled(true)
	return nil
}

func (n *sqliteLoadedDatabase) ApplyChanges(f *TMainForm, ndata *navData) (retErr error) {

	if len(ndata.bucketPath) != 2 {
		return errors.New("invalid selection")
	}
	tableName := ndata.bucketPath[1]

	// We have rendered row IDs, need to convert back to an SQLite primary key
	// TODO stash the real key inside f.contentBox.Objects()
	// FIXME breaks if you try and edit the primary key(!)

	ctx := context.Background()

	tx, err := n.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}

	var commitOK bool = false
	defer func() {
		if !commitOK {
			err := tx.Rollback()
			if err != nil {
				retErr = err
			}
		}
	}()

	// Data grid properties
	var columnNames []string
	for i := int32(0); i < f.contentBox.ColCount(); i++ {
		columnNames = append(columnNames, f.contentBox.Columns().Items(i).Title().Caption())
	}

	// Query sqlite table metadata to determine which of these is the PRIMARY KEY
	var primaryColumnName string
	err = tx.QueryRowContext(ctx, `SELECT l.name FROM pragma_table_info(?) as l WHERE l.pk = 1;`, tableName).Scan(&primaryColumnName)
	if err != nil {
		return fmt.Errorf("Finding primary key for update: %w", err)
	}

	// Convert it to an index
	var primaryColumnIdx int = -1
	for i := 0; i < len(columnNames); i++ {
		if columnNames[i] == primaryColumnName {
			primaryColumnIdx = i
			break
		}
	}
	if primaryColumnIdx == -1 {
		return fmt.Errorf("Primary key %q missing from available columns", primaryColumnName)
	}

	// SQLite can only LIMIT 1 on update/delete if it was compiled with
	// SQLITE_ENABLE_UPDATE_DELETE_LIMIT, which isn't the case for the mattn
	// cgo library
	// Skip that, and just rely on primary key uniqueness

	// Edit
	for rowid, editcells := range f.updateRows {
		stmt := `UPDATE "` + tableName + `" SET `
		params := []interface{}{} // FIXME reinstate types for the driver (although SQLite doesn't mind)

		for ct, cell := range editcells {
			if ct > 0 {
				stmt += `, `
			}
			stmt += `"` + columnNames[cell] + `" = ?`
			params = append(params, f.contentBox.Cells(cell, rowid))
		}
		stmt += ` WHERE "` + primaryColumnName + `" = ?`
		pkVal := f.contentBox.Cells(int32(primaryColumnIdx), rowid)
		params = append(params, pkVal)

		_, err = tx.ExecContext(ctx, stmt, params...)
		if err != nil {
			return fmt.Errorf("Updating row %q: %w", pkVal, err)
		}
	}

	// Delete by key (affects rowids after re-render)
	for rowid, _ := range f.deleteRows {
		pkVal := f.contentBox.Cells(int32(primaryColumnIdx), rowid)

		stmt := `DELETE FROM "` + tableName + `" WHERE "` + primaryColumnName + `" = ?`

		_, err = tx.ExecContext(ctx, stmt, pkVal)
		if err != nil {
			return fmt.Errorf("Deleting row %q: %w", pkVal, err)
		}
	}

	// Insert all new entries
	for rowid, _ := range f.insertRows {
		stmt := `INSERT INTO "` + tableName + `" (` + strings.Join(columnNames, `, `) + `) VALUES (`
		params := []interface{}{} // FIXME reinstate types for the driver (although SQLite doesn't mind)

		for colid := 0; colid < len(columnNames); colid++ {
			if colid > 0 {
				stmt += `, `
			}
			stmt += "?"
			params = append(params, f.contentBox.Cells(int32(colid), rowid))
		}
		stmt += `)`

		_, err = tx.ExecContext(ctx, stmt, params...)
		if err != nil {
			return fmt.Errorf("Inserting row: %w", err)
		}
	}

	err = tx.Commit()
	if err != nil {
		return err
	}

	commitOK = true // no need for rollback
	return nil
}

func (ld *sqliteLoadedDatabase) NavChildren(ndata *navData) ([]string, error) {

	if len(ndata.bucketPath) == 0 {
		// The top-level children are always:
		return []string{sqliteTablesCaption}, nil
	}

	if len(ndata.bucketPath) == 1 && ndata.bucketPath[0] == sqliteTablesCaption {
		rr, err := ld.db.Query(`SELECT name FROM sqlite_master WHERE type='table' ORDER BY name ASC;`)
		if err != nil {
			return nil, err
		}
		defer rr.Close()

		var gather []string
		for rr.Next() {
			var tableName string
			err = rr.Scan(&tableName)
			if err != nil {
				return nil, err
			}

			gather = append(gather, tableName)
		}
		if rr.Err() != nil {
			return nil, rr.Err()
		}

		return gather, nil
	}

	if len(ndata.bucketPath) == 2 {
		return nil, nil // Never any deeper children
	}

	return nil, fmt.Errorf("unknown nav path %#v", ndata.bucketPath)
}

func (ld *sqliteLoadedDatabase) NavContext(ndata *navData) (ret []contextAction, err error) {

	if len(ndata.bucketPath) == 0 {
		ret = append(ret, contextAction{"Compact database", ld.CompactDatabase})
		ret = append(ret, contextAction{"Export backup...", ld.ExportBackup})
	}
	if len(ndata.bucketPath) == 2 {
		ret = append(ret, contextAction{"Drop table", ld.DropTable})
	}
	return
}

func (ld *sqliteLoadedDatabase) CompactDatabase(sender vcl.IComponent, ndata *navData) error {
	_, err := ld.db.Exec(`VACUUM;`)
	return err
}

func (ld *sqliteLoadedDatabase) ExportBackup(sender vcl.IComponent, ndata *navData) error {
	// Popup for output file
	dlg := vcl.NewSaveDialog(sender)
	dlg.SetTitle("Save backup as...")
	dlg.SetFilter(sqliteFilter)
	ret := dlg.Execute() // Fake blocking
	if !ret {
		return nil // cancelled
	}

	_, err := ld.db.Exec(`VACUUM INTO ?`, dlg.FileName())
	return err
}

func (ld *sqliteLoadedDatabase) DropTable(sender vcl.IComponent, ndata *navData) error {
	if len(ndata.bucketPath) != 2 {
		return errors.New("Invalid selection")
	}

	//

	tableName := ndata.bucketPath[1]

	if !vcl_confirm_dialog(sender, "Drop table", fmt.Sprintf("Are you sure you want to drop the table %q?", tableName)) {
		return nil // cancelled
	}

	_, err := ld.db.Exec(`DROP TABLE "` + tableName + `"`) // WARNING can't prepare this parameter, but it comes from the DB (trusted)
	return err
}

func (ld *sqliteLoadedDatabase) Close() {
	_ = ld.db.Close()
	ld.arena = nil
}

var _ loadedDatabase = &sqliteLoadedDatabase{} // interface assertion

//

func (f *TMainForm) sqliteAddDatabaseFromFile(path string, cliDriver bool) {

	driver := "sqlite3"
	if cliDriver {
		driver = "sqliteclidriver"
	}

	// TODO load in background thread to stop blocking the UI
	db, err := sql.Open(driver, path)
	if err != nil {
		vcl.ShowMessage(fmt.Sprintf("Failed to load database '%s': %s", path, err.Error()))
		return
	}

	ld := &sqliteLoadedDatabase{
		path:        path,
		displayName: filepath.Base(path),
		db:          db,
	}

	ld.nav = f.Buckets.Items().Add(nil, ld.displayName)
	ld.nav.SetImageIndex(imgDatabase)
	ld.nav.SetSelectedIndex(imgDatabase)
	ld.nav.SetHasChildren(true) // dynamically populate in OnNavExpanding
	navData := &navData{
		ld:             ld,
		childrenLoaded: false,      // will be loaded dynamically
		bucketPath:     []string{}, // empty = root
	}
	ld.nav.SetData(unsafe.Pointer(navData))

	f.dbs = append(f.dbs, ld)
	f.Buckets.SetSelected(ld.nav) // Select new element

	ld.Keepalive(navData)
}