Merge #150350 #150388

150350: catalog: allow offline descriptors for type hydrations r=fqazi a=fqazi

Previously, when a function referenced an implicit table record type, the entire schema would become inaccessible if the referenced table was executing an IMPORT. This would happen because looking up a schema involves hydrating all type references, which include implicit table record types.

To address this, this patch allows for offline descriptors during the type hydration logic, since this scenario is possible during an import. Additionally, a test-only assertion is added to confirm this.

Fixes: #149988
Release note (bug fix): Fixed a bug where the entire schema would become inaccessible if a table was referenced as an implicit record type by a UDF while the table was undergoing an IMPORT.

150388: util/parquet,sql: clean up DOidWrapper handling a bit r=yuzefovich a=yuzefovich

This commit removes redundant branch for `DOidWrapper` in `encodeArrayElement` (we already do the unwrapping at the top of the function, so this was dead code) as well as replaces the custom unwrap function in `parquet` package with the exported one. It also replaces some other unwrappers with that exported function.

Epic: None
Release note: None

Co-authored-by: Faizan Qazi <[email protected]>
Co-authored-by: Yahor Yuzefovich <[email protected]>

Author: 3 people

pkg/sql/catalog/descs/BUILD.bazel

@@ -65,6 +65,7 @@ go_library(
         "//pkg/sql/sqlerrors",
         "//pkg/sql/sqlliveness",
         "//pkg/sql/types",
+        "//pkg/util/buildutil",
         "//pkg/util/hlc",
         "//pkg/util/intsets",
         "//pkg/util/iterutil",

pkg/sql/catalog/descs/descriptor.go

@@ -215,15 +215,17 @@ func getDescriptorsByID(
 	if err := tc.finalizeDescriptors(ctx, txn, flags, descs, vls); err != nil {
 		return err
 	}
-	// Hydration is skipped if "SkipHydration" flag is true.
-	if err := tc.hydrateDescriptors(ctx, txn, flags, descs); err != nil {
-		return err
-	}
+	// Apply any filters on descriptors before hydrating, since if a descriptor
+	// is offline / dropped, we are doing needless work.
 	for _, desc := range descs {
 		if err := filterDescriptor(desc, flags); err != nil {
 			return err
 		}
 	}
+	// Hydration is skipped if "SkipHydration" flag is true.
+	if err := tc.hydrateDescriptors(ctx, txn, flags, descs); err != nil {
+		return err
+	}
 	return nil
 }
 

pkg/sql/catalog/descs/hydrate.go

@@ -14,10 +14,12 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/nstree"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/schemadesc"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/intsets"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/errors"
@@ -67,7 +69,7 @@ func (tc *Collection) hydrateDescriptors(
 
 	// hydrate mutable hydratable descriptors of the slice in-place.
 	if !hydratableMutableIndexes.Empty() {
-		typeFn := makeMutableTypeLookupFunc(tc, txn, descs)
+		typeFn := makeMutableTypeLookupFunc(tc, txn, flags, descs)
 		for _, i := range hydratableMutableIndexes.Ordered() {
 			if err := hydrate(ctx, descs[i], typeFn); err != nil {
 				return err
@@ -108,7 +110,7 @@ func (tc *Collection) hydrateDescriptors(
 }
 
 func makeMutableTypeLookupFunc(
-	tc *Collection, txn *kv.Txn, descs []catalog.Descriptor,
+	tc *Collection, txn *kv.Txn, flags getterFlags, descs []catalog.Descriptor,
 ) typedesc.TypeLookupFunc {
 	var mc nstree.MutableCatalog
 	for _, desc := range descs {
@@ -131,7 +133,7 @@ func makeMutableTypeLookupFunc(
 		if id == catconstants.PublicSchemaID {
 			return schemadesc.GetPublicSchema(), nil
 		}
-		flags := getterFlags{
+		f := getterFlags{
 			contextFlags: contextFlags{
 				isMutable: true,
 			},
@@ -140,9 +142,35 @@ func makeMutableTypeLookupFunc(
 				withoutLeased:    true,
 				withoutHydration: skipHydration,
 			},
+			descFilters: descFilters{
+				// For hydration, we will allow offline descriptors for lookups
+				// of type descriptors. A descriptor can be offline either due to:
+				// 1) IMPORT in which case we are looking at an implicit record type
+				// 2) RESTORE which will always pick a new object to restore into, so
+				//    only the restore code paths can enter here, which will include
+				//    offline descriptors explicitly.
+				withoutOffline: false,
+			},
+		}
+		g := ByIDGetter(makeGetterBase(txn, tc, f))
+		desc, err := g.Desc(ctx, id)
+		if err != nil {
+			return nil, err
 		}
-		g := ByIDGetter(makeGetterBase(txn, tc, flags))
-		return g.Desc(ctx, id)
+		// Sanity: If offline descriptors were asked to be ignored, then we should
+		// only observe ones caused by IMPORT.
+		if flags.descFilters.withoutOffline &&
+			(desc.GetOfflineReason() != "" && desc.GetOfflineReason() != tabledesc.OfflineReasonImporting) {
+			if buildutil.CrdbTestBuild {
+				return nil, errors.AssertionFailedf("unexpected offline descriptor %s(%d): %s",
+					desc.GetName(),
+					desc.GetID(),
+					desc.GetOfflineReason())
+			}
+			// For release builds surface an offline descriptor error.
+			return nil, catalog.FilterOfflineDescriptor(desc)
+		}
+		return desc, nil
 	}
 	return makeTypeLookupFuncForHydration(mc, mutableLookupFunc)
 }
@@ -169,11 +197,34 @@ func makeImmutableTypeLookupFunc(
 			},
 			descFilters: descFilters{
 				withoutDropped: true,
-				withoutOffline: flags.descFilters.withoutOffline,
+				// For hydration, we will allow offline descriptors for lookups
+				// of type descriptors. A descriptor can be offline either due to:
+				// 1) IMPORT in which case we are looking at an implicit record type
+				// 2) RESTORE which will always pick a new object to restore into, so
+				//    only the restore code paths can enter here, which will include
+				//    offline descriptors explicitly.
+				withoutOffline: false,
 			},
 		}
 		g := ByIDGetter(makeGetterBase(txn, tc, f))
-		return g.Desc(ctx, id)
+		desc, err := g.Desc(ctx, id)
+		if err != nil {
+			return nil, err
+		}
+		// Sanity: If offline descriptors were asked to be ignored, then we should
+		// only observe ones caused by IMPORT.
+		if flags.descFilters.withoutOffline &&
+			(desc.GetOfflineReason() != "" && desc.GetOfflineReason() != tabledesc.OfflineReasonImporting) {
+			if buildutil.CrdbTestBuild {
+				return nil, errors.AssertionFailedf("unexpected offline descriptor %s(%d): %s",
+					desc.GetName(),
+					desc.GetID(),
+					desc.GetOfflineReason())
+			}
+			// For release builds surface an offline descriptor error.
+			return nil, catalog.FilterOfflineDescriptor(desc)
+		}
+		return desc, nil
 	}
 	return makeTypeLookupFuncForHydration(mc, immutableLookupFunc)
 }

pkg/sql/catalog/resolver/resolver_test.go

@@ -764,6 +764,64 @@ CREATE INDEX baz_idx ON baz (s);
 	require.NoError(t, err)
 }
 
+// TestOfflineImplicitTableRecord confirms that offline implicit record references
+// do not impact unrelated objects.
+func TestOfflineImplicitTableRecord(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	srv, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{})
+	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
+	tDB := sqlutils.MakeSQLRunner(sqlDB)
+
+	tDB.Exec(t, `
+CREATE DATABASE db1;
+USE db1;
+CREATE TABLE a (k INT PRIMARY KEY);
+
+CREATE TABLE b (
+  k INT PRIMARY KEY,
+  a INT,
+  b INT
+);
+
+-- This function will have an implicit record reference to b.
+CREATE FUNCTION f(x INT) RETURNS b LANGUAGE SQL AS $$
+  SELECT k, a, b FROM b WHERE k = x
+$$;
+`)
+	// Mark the function as offline.
+	var tblID descpb.ID
+	err := sqltestutils.TestingDescsTxn(ctx, s, func(ctx context.Context, txn isql.Txn, col *descs.Collection) error {
+		dbID, err := col.LookupDatabaseID(ctx, txn.KV(), "db1")
+		if err != nil {
+			return err
+		}
+		schemaID, err := col.LookupSchemaID(ctx, txn.KV(), dbID, "public")
+		if err != nil {
+			return err
+		}
+		tblID, err = col.LookupObjectID(ctx, txn.KV(), dbID, schemaID, "b")
+		if err != nil {
+			return err
+		}
+		mutTbl, err := col.MutableByID(txn.KV()).Table(ctx, tblID)
+		if err != nil {
+			return err
+		}
+		mutTbl.SetOffline(tabledesc.OfflineReasonImporting)
+		return col.WriteDesc(ctx, false /*kvTrace*/, mutTbl, txn.KV())
+	})
+	require.NoError(t, err)
+	// Confirm that we can still resolve things under the schema.
+	tDB.Exec(t, "SELECT * FROM db1.public.a")
+	// Confirm resolving the actual table hits an error.
+	tDB.ExpectErr(t, "relation \"b\" is offline: importing", "SELECT * FROM db1.public.b")
+	tDB.ExpectErr(t, "relation \"b\" is offline: importing", "SELECT * FROM f(0)")
+}
+
 func newTableIndexName(db, sc, tbl, idx string) *tree.TableIndexName {
 	name := &tree.TableIndexName{
 		Index: tree.UnrestrictedName(idx),

pkg/sql/colconv/datum_to_vec.eg.go

@@ -69,10 +69,7 @@ type familyWidthPair struct {
 
 var rowToVecPreludeTmpls = map[familyWidthPair]string{
 	{types.StringFamily, anyWidth}: `// Handle other STRING-related OID types, like oid.T_name.
-												 wrapper, ok := %[1]s.(*tree.DOidWrapper)
-												 if ok {
-												     %[1]s = wrapper.Wrapped
-												 }`,
+												 %[1]s = tree.UnwrapDOidWrapper(%[1]s)`,
 }
 
 // rowToVecConversionTmpls maps the type families to the corresponding

pkg/sql/colexec/execgen/cmd/execgen/rowtovec_gen.go

@@ -245,9 +245,7 @@ func (c *CustomFuncs) ConvertConstArrayToTuple(arr *memo.ConstExpr) opt.ScalarEx
 // collated string constant with the given locale.
 func (c *CustomFuncs) CastToCollatedString(str opt.ScalarExpr, locale string) opt.ScalarExpr {
 	datum := str.(*memo.ConstExpr).Value
-	if wrap, ok := datum.(*tree.DOidWrapper); ok {
-		datum = wrap.Wrapped
-	}
+	datum = tree.UnwrapDOidWrapper(datum)
 
 	// buildCollated is a recursive helper function to handle casting arrays.
 	var buildCollated func(datum tree.Datum) tree.Datum

pkg/sql/colexec/rowtovec.eg.go

@@ -336,8 +336,6 @@ func encodeArrayElement(b []byte, d tree.Datum) ([]byte, error) {
 		return encoding.EncodeUntaggedIntValue(b, int64(t.Oid)), nil
 	case *tree.DCollatedString:
 		return encoding.EncodeUntaggedBytesValue(b, []byte(t.Contents)), nil
-	case *tree.DOidWrapper:
-		return encodeArrayElement(b, t.Wrapped)
 	case *tree.DEnum:
 		return encoding.EncodeUntaggedBytesValue(b, t.PhysicalRep), nil
 	case *tree.DJSON:

pkg/sql/opt/norm/scalar_funcs.go

@@ -903,7 +903,7 @@ func TestAnalyzeSystemTables(t *testing.T) {
 		return descpb.ID(tableID)
 	}
 	for _, row := range rows {
-		tableName := string(*row[0].(*tree.DOidWrapper).Wrapped.(*tree.DString))
+		tableName := string(*tree.UnwrapDOidWrapper(row[0]).(*tree.DString))
 		if DisallowedOnSystemTable(getTableID(tableName)) {
 			continue
 		}