importccl: add support for IMPORT INTO RBR table

pkg/ccl/importccl/import_stmt.go

@@ -974,6 +974,24 @@ func importPlanHook(
 			}
 		}
 
+		// Store the primary region of the database being imported into. This is
+		// used during job execution to evaluate certain default expressions and
+		// computed columns such as `gateway_region`.
+		var databasePrimaryRegion descpb.RegionName
+		if db.IsMultiRegion() {
+			if err := sql.DescsTxn(ctx, p.ExecCfg(), func(ctx context.Context, txn *kv.Txn,
+				descsCol *descs.Collection) error {
+				regionConfig, err := sql.SynthesizeRegionConfig(ctx, txn, db.GetID(), descsCol)
+				if err != nil {
+					return err
+				}
+				databasePrimaryRegion = regionConfig.PrimaryRegion()
+				return nil
+			}); err != nil {
+				return errors.Wrap(err, "failed to resolve region config for multi region database")
+			}
+		}
+
 		telemetry.CountBucketed("import.files", int64(len(files)))
 
 		// Record telemetry for userfile being used as the import target.
@@ -1002,16 +1020,17 @@ func importPlanHook(
 		// StartableJob which we attached to the connExecutor somehow.
 
 		importDetails := jobspb.ImportDetails{
-			URIs:              files,
-			Format:            format,
-			ParentID:          db.GetID(),
-			Tables:            tableDetails,
-			Types:             typeDetails,
-			SSTSize:           sstSize,
-			Oversample:        oversample,
-			SkipFKs:           skipFKs,
-			ParseBundleSchema: importStmt.Bundle,
-			DefaultIntSize:    p.SessionData().DefaultIntSize,
+			URIs:                  files,
+			Format:                format,
+			ParentID:              db.GetID(),
+			Tables:                tableDetails,
+			Types:                 typeDetails,
+			SSTSize:               sstSize,
+			Oversample:            oversample,
+			SkipFKs:               skipFKs,
+			ParseBundleSchema:     importStmt.Bundle,
+			DefaultIntSize:        p.SessionData().DefaultIntSize,
+			DatabasePrimaryRegion: databasePrimaryRegion,
 		}
 
 		jr := jobs.Record{
@@ -1295,9 +1314,6 @@ func prepareExistingTableDescForIngestion(
 	if len(desc.Mutations) > 0 {
 		return nil, errors.Errorf("cannot IMPORT INTO a table with schema changes in progress -- try again later (pending mutation %s)", desc.Mutations[0].String())
 	}
-	if desc.LocalityConfig != nil && desc.LocalityConfig.GetRegionalByRow() != nil {
-		return nil, unimplemented.NewWithIssueDetailf(61133, "import.regional-by-row", "IMPORT into REGIONAL BY ROW table not supported")
-	}
 
 	// Note that desc is just used to verify that the version matches.
 	importing, err := descsCol.GetMutableTableVersionByID(ctx, desc.ID, txn)

pkg/ccl/importccl/import_stmt_test.go

@@ -6429,6 +6429,14 @@ func TestImportMultiRegion(t *testing.T) {
 
 	simpleOcf := fmt.Sprintf("nodelocal://0/avro/%s", "simple.ocf")
 
+	var data string
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method == "GET" {
+			_, _ = w.Write([]byte(data))
+		}
+	}))
+	defer srv.Close()
+
 	// Table schemas for USING
 	tableSchemaMR := fmt.Sprintf("nodelocal://0/avro/%s", "simple-schema-multi-region.sql")
 	tableSchemaMRRegionalByRow := fmt.Sprintf("nodelocal://0/avro/%s",
@@ -6498,6 +6506,7 @@ DROP VIEW IF EXISTS v`,
 			create    string
 			args      []interface{}
 			errString string
+			data      string
 		}{
 			{
 				name:      "import-create-using-multi-region-to-non-multi-region-database",
@@ -6523,16 +6532,34 @@ DROP VIEW IF EXISTS v`,
 				errString: "IMPORT to REGIONAL BY ROW table not supported",
 			},
 			{
-				name:      "import-into-multi-region-regional-by-row-to-multi-region-database",
+				name:   "import-into-multi-region-regional-by-row-default-col-to-multi-region-database",
+				db:     "multi_region",
+				table:  "mr_regional_by_row",
+				create: "CREATE TABLE mr_regional_by_row (i INT8 PRIMARY KEY, s text, b bytea) LOCALITY REGIONAL BY ROW",
+				sql:    "IMPORT INTO mr_regional_by_row AVRO DATA ($1)",
+				args:   []interface{}{simpleOcf},
+			},
+			{
+				name:   "import-into-multi-region-regional-by-row-to-multi-region-database",
+				db:     "multi_region",
+				table:  "mr_regional_by_row",
+				create: "CREATE TABLE mr_regional_by_row (i INT8 PRIMARY KEY, s text, b bytea) LOCALITY REGIONAL BY ROW",
+				sql:    "IMPORT INTO mr_regional_by_row (i, s, b, crdb_region) CSV DATA ($1)",
+				args:   []interface{}{srv.URL},
+				data:   "1,\"foo\",NULL,us-east1\n",
+			},
+			{
+				name:      "import-into-multi-region-regional-by-row-to-multi-region-database-wrong-value",
 				db:        "multi_region",
 				table:     "mr_regional_by_row",
 				create:    "CREATE TABLE mr_regional_by_row (i INT8 PRIMARY KEY, s text, b bytea) LOCALITY REGIONAL BY ROW",
-				sql:       "IMPORT INTO mr_regional_by_row AVRO DATA ($1)",
-				args:      []interface{}{simpleOcf},
-				errString: "IMPORT into REGIONAL BY ROW table not supported",
+				sql:       "IMPORT INTO mr_regional_by_row (i, s, b, crdb_region) CSV DATA ($1)",
+				args:      []interface{}{srv.URL},
+				data:      "1,\"foo\",NULL,us-west1\n",
+				errString: "invalid input value for enum crdb_internal_region",
 			},
 			{
-				name:   "import-into-using-multi-region-global-to-multi-region-database",
+				name:   "import-into-multi-region-global-to-multi-region-database",
 				db:     "multi_region",
 				table:  "mr_global",
 				create: "CREATE TABLE mr_global (i INT8 PRIMARY KEY, s text, b bytea) LOCALITY GLOBAL",
@@ -6549,14 +6576,18 @@ DROP VIEW IF EXISTS v`,
 				_, err = sqlDB.Exec(fmt.Sprintf("DROP TABLE IF EXISTS %q CASCADE", test.table))
 				require.NoError(t, err)
 
+				if test.data != "" {
+					data = test.data
+				}
+
 				if test.create != "" {
 					_, err = sqlDB.Exec(test.create)
 					require.NoError(t, err)
 				}
 
 				_, err = sqlDB.ExecContext(context.Background(), test.sql, test.args...)
 				if test.errString != "" {
-					testutils.IsError(err, test.errString)
+					require.True(t, testutils.IsError(err, test.errString))
 				} else {
 					require.NoError(t, err)
 					res := sqlDB.QueryRow(fmt.Sprintf("SELECT count(*) FROM %q", test.table))

pkg/ccl/importccl/import_table_creation.go

@@ -205,7 +205,7 @@ func MakeSimpleTableDescriptor(
 	evalCtx := tree.EvalContext{
 		Context:            ctx,
 		Sequence:           &importSequenceOperators{},
-		Regions:            &importRegionOperator{},
+		Regions:            makeImportRegionOperator(""),
 		SessionDataStack:   sessiondata.NewStack(&sessiondata.SessionData{}),
 		ClientNoticeSender: &faketreeeval.DummyClientNoticeSender{},
 		Settings:           st,
@@ -264,13 +264,38 @@ var (
 )
 
 // Implements the tree.RegionOperator interface.
-type importRegionOperator struct{}
+type importRegionOperator struct {
+	primaryRegion descpb.RegionName
+}
+
+func makeImportRegionOperator(primaryRegion descpb.RegionName) *importRegionOperator {
+	return &importRegionOperator{primaryRegion: primaryRegion}
+}
+
+// importDatabaseRegionConfig is a stripped down version of
+// multiregion.RegionConfig that is used by import.
+type importDatabaseRegionConfig struct {
+	primaryRegion descpb.RegionName
+}
+
+// IsValidRegionNameString implements the tree.DatabaseRegionConfig interface.
+func (i importDatabaseRegionConfig) IsValidRegionNameString(_ string) bool {
+	// Unimplemented.
+	return false
+}
+
+// PrimaryRegionString implements the tree.DatabaseRegionConfig interface.
+func (i importDatabaseRegionConfig) PrimaryRegionString() string {
+	return string(i.primaryRegion)
+}
+
+var _ tree.DatabaseRegionConfig = &importDatabaseRegionConfig{}
 
 // CurrentDatabaseRegionConfig is part of the tree.EvalDatabase interface.
 func (so *importRegionOperator) CurrentDatabaseRegionConfig(
 	_ context.Context,
 ) (tree.DatabaseRegionConfig, error) {
-	return nil, errors.WithStack(errRegionOperator)
+	return importDatabaseRegionConfig{primaryRegion: so.primaryRegion}, nil
 }
 
 // ValidateAllMultiRegionZoneConfigsInCurrentDatabase is part of the tree.EvalDatabase interface.

pkg/ccl/importccl/read_import_base.go

@@ -62,6 +62,7 @@ func runImport(
 	// TODO(adityamaru): Should we just plumb the flowCtx instead of this
 	// assignment.
 	evalCtx.DB = flowCtx.Cfg.DB
+	evalCtx.Regions = makeImportRegionOperator(spec.DatabasePrimaryRegion)
 	semaCtx := tree.MakeSemaContext()
 	semaCtx.TypeResolver = importResolver
 	conv, err := makeInputConverter(ctx, &semaCtx, spec, evalCtx, kvCh, seqChunkProvider)