Merge pull request #1961 from dolthub/aaron/readonly-engine-updates

engine: Make ReadOnly a toggleable atomic.Bool.

Author: reltuk

engine.go

@@ -20,6 +20,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/dolthub/vitess/go/sqltypes"
 	querypb "github.com/dolthub/vitess/go/vt/proto/query"
@@ -136,7 +137,7 @@ type Engine struct {
 	ProcessList       sql.ProcessList
 	MemoryManager     *sql.MemoryManager
 	BackgroundThreads *sql.BackgroundThreads
-	IsReadOnly        bool
+	ReadOnly          atomic.Bool
 	IsServerLocked    bool
 	PreparedDataCache *PreparedDataCache
 	mu                *sync.Mutex
@@ -168,17 +169,18 @@ func New(a *analyzer.Analyzer, cfg *Config) *Engine {
 	})
 	a.Catalog.RegisterFunction(emptyCtx, function.GetLockingFuncs(ls)...)
 
-	return &Engine{
+	ret := &Engine{
 		Analyzer:          a,
 		MemoryManager:     sql.NewMemoryManager(sql.ProcessMemory),
 		ProcessList:       NewProcessList(),
 		LS:                ls,
 		BackgroundThreads: sql.NewBackgroundThreads(),
-		IsReadOnly:        cfg.IsReadOnly,
 		IsServerLocked:    cfg.IsServerLocked,
 		PreparedDataCache: NewPreparedDataCache(),
 		mu:                &sync.Mutex{},
 	}
+	ret.ReadOnly.Store(cfg.IsReadOnly)
+	return ret
 }
 
 // NewDefault creates a new default Engine.
@@ -635,23 +637,19 @@ func (e *Engine) WithBackgroundThreads(b *sql.BackgroundThreads) *Engine {
 	return e
 }
 
+func (e *Engine) IsReadOnly() bool {
+	return e.ReadOnly.Load()
+}
+
 // readOnlyCheck checks to see if the query is valid with the modification setting of the engine.
 func (e *Engine) readOnlyCheck(node sql.Node) error {
-	if plan.IsDDLNode(node) {
-		if e.IsReadOnly {
-			return sql.ErrReadOnly.New()
-		} else if e.IsServerLocked {
-			return sql.ErrDatabaseWriteLocked.New()
-		}
+	// Note: We only compute plan.IsReadOnly if the server is in one of
+	// these two modes, since otherwise it is simply wasted work.
+	if e.IsReadOnly() && !plan.IsReadOnly(node) {
+		return sql.ErrReadOnly.New()
 	}
-	switch node.(type) {
-	case
-		*plan.DeleteFrom, *plan.InsertInto, *plan.Update, *plan.LockTables, *plan.UnlockTables:
-		if e.IsReadOnly {
-			return sql.ErrReadOnly.New()
-		} else if e.IsServerLocked {
-			return sql.ErrDatabaseWriteLocked.New()
-		}
+	if e.IsServerLocked && !plan.IsReadOnly(node) {
+		return sql.ErrDatabaseWriteLocked.New()
 	}
 	return nil
 }

enginetest/enginetests.go

@@ -623,7 +623,7 @@ func TestOrderByGroupBy(t *testing.T, harness Harness) {
 func TestReadOnly(t *testing.T, harness Harness) {
 	harness.Setup(setup.Mytable...)
 	e := mustNewEngine(t, harness)
-	e.IsReadOnly = true
+	e.ReadOnly.Store(true)
 	defer e.Close()
 
 	RunQuery(t, e, harness, `SELECT i FROM mytable`)
@@ -634,6 +634,11 @@ func TestReadOnly(t *testing.T, harness Harness) {
 		`INSERT INTO mytable (i, s) VALUES(42, 'yolo')`,
 		`CREATE VIEW myview3 AS SELECT i FROM mytable`,
 		`DROP VIEW myview`,
+		`DROP DATABASE mydb`,
+		`CREATE DATABASE newdb`,
+		`CREATE USER tester@localhost`,
+		`CREATE ROLE test_role`,
+		`GRANT SUPER ON * TO 'root'@'localhost'`,
 	}
 
 	for _, query := range writingQueries {

sql/plan/process.go

@@ -478,3 +478,49 @@ func IsShowNode(node sql.Node) bool {
 func IsNoRowNode(node sql.Node) bool {
 	return IsDDLNode(node) || IsShowNode(node)
 }
+
+func IsReadOnly(node sql.Node) bool {
+	isExplain := false
+	transform.Inspect(node, func(n sql.Node) bool {
+		switch n.(type) {
+		case *DescribeQuery:
+			isExplain = true
+			return false
+		}
+		return true
+	})
+	if isExplain {
+		return true
+	}
+
+	if IsDDLNode(node) {
+		return false
+	}
+
+	switch node.(type) {
+	case *DeleteFrom, *InsertInto, *Update, *LockTables, *UnlockTables:
+		return false
+	}
+
+	isPrivNodeP := func(n sql.Node) bool {
+		switch node.(type) {
+		case *CreateUser, *DropUser, *RenameUser, *CreateRole, *DropRole, *Grant, *GrantRole, *GrantProxy, *Revoke, *RevokeRole, *RevokeAll, *RevokeProxy:
+			return true
+		}
+		return false
+	}
+	isPrivNode := false
+	transform.Inspect(node, func(n sql.Node) bool {
+		if isPrivNodeP(n) {
+			isPrivNode = true
+			return false
+		}
+		return true
+	})
+
+	if isPrivNode {
+		return false
+	}
+
+	return true
+}