Merge pull request #7539 from dolthub/fulghum/schema-pinning

Feature: Schema overriding

go/libraries/doltcore/doltdb/table.go

@@ -61,7 +61,8 @@ func IsValidIdentifier(name string) bool {
 
 // Table is a struct which holds row data, as well as a reference to its schema.
 type Table struct {
-	table durable.Table
+	table            durable.Table
+	overriddenSchema schema.Schema
 }
 
 // NewNomsTable creates a noms Struct which stores row data, index data, and schema.
@@ -115,6 +116,17 @@ func (t *Table) NodeStore() tree.NodeStore {
 	return durable.NodeStoreFromTable(t.table)
 }
 
+// OverrideSchema sets |sch| as the schema for this table, causing rows from this table to be transformed
+// into that schema when they are read from this table.
+func (t *Table) OverrideSchema(sch schema.Schema) {
+	t.overriddenSchema = sch
+}
+
+// GetOverriddenSchema returns the overridden schema if one has been set, otherwise it returns nil.
+func (t *Table) GetOverriddenSchema() schema.Schema {
+	return t.overriddenSchema
+}
+
 // SetConflicts sets the merge conflicts for this table.
 func (t *Table) SetConflicts(ctx context.Context, schemas conflict.ConflictSchema, conflictData durable.ConflictIndex) (*Table, error) {
 	table, err := t.table.SetConflicts(ctx, schemas, conflictData)

go/libraries/doltcore/sqle/database.go

@@ -254,16 +254,7 @@ func (db Database) GetTableInsensitive(ctx *sql.Context, tblName string) (sql.Ta
 		return nil, false, err
 	}
 
-	tbl, ok, err := db.getTableInsensitive(ctx, nil, ds, root, tblName)
-	if err != nil {
-		return nil, false, err
-	}
-
-	if !ok {
-		return nil, false, nil
-	}
-
-	return tbl, true, nil
+	return db.getTableInsensitive(ctx, nil, ds, root, tblName)
 }
 
 // GetTableInsensitiveAsOf implements sql.VersionedDatabase
@@ -294,17 +285,23 @@ func (db Database) GetTableInsensitiveAsOf(ctx *sql.Context, tableName string, a
 		return table, ok, nil
 	}
 
-	versionableTable, ok := table.(dtables.VersionableTable)
-	if !ok {
-		panic(fmt.Sprintf("unexpected table type %T", table))
-	}
+	switch t := table.(type) {
+	case dtables.VersionableTable:
+		versionedTable, err := t.LockedToRoot(ctx, root)
+		if err != nil {
+			return nil, false, err
+		}
+		return versionedTable, true, nil
 
-	versionedTable, err := versionableTable.LockedToRoot(ctx, root)
+	case *plan.EmptyTable:
+		// getTableInsensitive returns *plan.EmptyTable if the table doesn't exist in the data root, but
+		// schemas have been locked to a commit where the table does exist. Since the table is empty,
+		// there's no need to lock it to a root.
+		return t, true, nil
 
-	if err != nil {
-		return nil, false, err
+	default:
+		return nil, false, fmt.Errorf("unexpected table type %T", table)
 	}
-	return versionedTable, true, nil
 
 }
 
@@ -492,7 +489,17 @@ func (db Database) getTableInsensitive(ctx *sql.Context, head *doltdb.Commit, ds
 	}
 
 	// TODO: this should reuse the root, not lookup the db state again
-	return db.getTable(ctx, root, tblName)
+	table, found, err := db.getTable(ctx, root, tblName)
+	if err != nil {
+		return nil, false, err
+	}
+	if found {
+		return table, found, err
+	}
+
+	// If the table wasn't found in the specified data root, check if there is an overridden
+	// schema commit that contains it and return an empty table if so.
+	return resolveOverriddenNonexistentTable(ctx, tblName, db)
 }
 
 // resolveAsOf resolves given expression to a commit, if one exists.
@@ -636,14 +643,22 @@ func (db Database) getTable(ctx *sql.Context, root *doltdb.RootValue, tableName
 		return nil, false, fmt.Errorf("no state for database %s", db.RevisionQualifiedName())
 	}
 
-	key, err := doltdb.NewDataCacheKey(root)
+	overriddenSchemaRoot, err := resolveOverriddenSchemaRoot(ctx, db)
 	if err != nil {
 		return nil, false, err
 	}
 
-	cachedTable, ok := dbState.SessionCache().GetCachedTable(key, tableName)
-	if ok {
-		return cachedTable, true, nil
+	// If schema hasn't been overridden, we can use a cached table if one exists
+	if overriddenSchemaRoot == nil {
+		key, err := doltdb.NewDataCacheKey(root)
+		if err != nil {
+			return nil, false, err
+		}
+
+		cachedTable, ok := dbState.SessionCache().GetCachedTable(key, tableName)
+		if ok {
+			return cachedTable, true, nil
+		}
 	}
 
 	tableNames, err := getAllTableNames(ctx, root)
@@ -669,12 +684,26 @@ func (db Database) getTable(ctx *sql.Context, root *doltdb.RootValue, tableName
 		return nil, false, err
 	}
 
+	if overriddenSchemaRoot != nil {
+		err = overrideSchemaForTable(ctx, tableName, tbl, overriddenSchemaRoot)
+		if err != nil {
+			return nil, false, err
+		}
+	}
+
 	table, err := db.newDoltTable(tableName, sch, tbl)
 	if err != nil {
 		return nil, false, err
 	}
 
-	dbState.SessionCache().CacheTable(key, tableName, table)
+	// If the schema hasn't been overridden, cache the table
+	if overriddenSchemaRoot == nil {
+		key, err := doltdb.NewDataCacheKey(root)
+		if err != nil {
+			return nil, false, err
+		}
+		dbState.SessionCache().CacheTable(key, tableName, table)
+	}
 
 	return table, true, nil
 }

go/libraries/doltcore/sqle/database_provider.go

@@ -215,6 +215,26 @@ func (p *DoltDatabaseProvider) Database(ctx *sql.Context, name string) (sql.Data
 		return nil, sql.ErrDatabaseNotFound.New(name)
 	}
 
+	overriddenSchemaValue, err := getOverriddenSchemaValue(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	// If a schema override is set, ensure we're using a ReadOnlyDatabase
+	if overriddenSchemaValue != "" {
+		// TODO: It would be nice if we could set a "read-only reason" for the read only database and let people know
+		//       that the database is read-only because of the @@dolt_override_schema setting and that customers need
+		//       to unset that session variable to get a write query to work. Otherwise it may be confusing why a
+		//       write query isn't working.
+		if _, ok := database.(ReadOnlyDatabase); !ok {
+			readWriteDatabase, ok := database.(Database)
+			if !ok {
+				return nil, fmt.Errorf("expected an instance of sqle.Database, but found: %T", database)
+			}
+			return ReadOnlyDatabase{readWriteDatabase}, nil
+		}
+	}
+
 	return database, nil
 }
 

go/libraries/doltcore/sqle/dsess/variables.go

@@ -38,6 +38,7 @@ const (
 	TransactionsDisabledSysVar           = "dolt_transactions_disabled"
 	ForceTransactionCommit               = "dolt_force_transaction_commit"
 	CurrentBatchModeKey                  = "batch_mode"
+	DoltOverrideSchema                   = "dolt_override_schema"
 	AllowCommitConflicts                 = "dolt_allow_commit_conflicts"
 	ReplicateToRemote                    = "dolt_replicate_to_remote"
 	ReadReplicaRemote                    = "dolt_read_replica_remote"

go/libraries/doltcore/sqle/enginetest/dolt_engine_test.go

@@ -116,6 +116,29 @@ func TestSingleQuery(t *testing.T) {
 	enginetest.TestQueryWithEngine(t, harness, engine, test)
 }
 
+func TestSchemaOverrides(t *testing.T) {
+	tcc := &testCommitClock{}
+	cleanup := installTestCommitClock(tcc)
+	defer cleanup()
+
+	for _, script := range SchemaOverrideTests {
+		sql.RunWithNowFunc(tcc.Now, func() error {
+			harness := newDoltHarness(t)
+			harness.Setup(setup.MydbData)
+
+			engine, err := harness.NewEngine(t)
+			if err != nil {
+				panic(err)
+			}
+			// engine.EngineAnalyzer().Debug = true
+			// engine.EngineAnalyzer().Verbose = true
+
+			enginetest.TestScriptWithEngine(t, engine, harness, script)
+			return nil
+		})
+	}
+}
+
 // Convenience test for debugging a single query. Unskip and set to the desired query.
 func TestSingleScript(t *testing.T) {
 	t.Skip()