From c7932c56779690390f8d7169bf58c1ffa9518bd9 Mon Sep 17 00:00:00 2001
From: Adam Storm <storm@cockroachlabs.com>
Date: Mon, 22 Feb 2021 10:36:24 -0500
Subject: [PATCH] sql: Enable IMPORT of tables into multi-region databases

Previously, tables that were exported from non-multi-region databases
were not able to be imported into multi-region databases. This commit
enables the above operation.

Release note: None

Release justification: Bug in interaction with existing feature and new
feature.
---
 pkg/ccl/backupccl/restore_job.go              |  18 ++++
 .../logic_test/multi_region_import_export     | 100 ++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100644 pkg/ccl/logictestccl/testdata/logic_test/multi_region_import_export

diff --git a/pkg/ccl/backupccl/restore_job.go b/pkg/ccl/backupccl/restore_job.go
index 955459c130d4..417ff8d5cc44 100644
--- a/pkg/ccl/backupccl/restore_job.go
+++ b/pkg/ccl/backupccl/restore_job.go
@@ -363,6 +363,24 @@ func WriteDescriptors(
 						table.GetID(), table)
 				}
 			}
+
+			// If the table descriptor is being written to a multi-region database and
+			// the table does not have a locality config setup, set one up here. The
+			// table's locality config will be set to the default locality - REGIONAL
+			// BY TABLE IN PRIMARY REGION.
+			_, dbDesc, err := descsCol.GetImmutableDatabaseByID(
+				ctx, txn, table.GetParentID(), tree.DatabaseLookupFlags{
+					Required:       true,
+					AvoidCached:    true,
+					IncludeOffline: true,
+				})
+			if err != nil {
+				return err
+			}
+			if dbDesc.GetRegionConfig() != nil && table.GetLocalityConfig() == nil {
+				table.(*tabledesc.Mutable).SetTableLocalityRegionalByTable(tree.PrimaryRegionLocalityName)
+			}
+
 			if err := descsCol.WriteDescToBatch(
 				ctx, false /* kvTrace */, tables[i].(catalog.MutableDescriptor), b,
 			); err != nil {
diff --git a/pkg/ccl/logictestccl/testdata/logic_test/multi_region_import_export b/pkg/ccl/logictestccl/testdata/logic_test/multi_region_import_export
new file mode 100644
index 000000000000..b806c9416af6
--- /dev/null
+++ b/pkg/ccl/logictestccl/testdata/logic_test/multi_region_import_export
@@ -0,0 +1,100 @@
+# LogicTest: multiregion-9node-3region-3azs
+
+query TTTT colnames
+SHOW REGIONS
+----
+region          zones                   database_names  primary_region_of
+ap-southeast-2  {ap-az1,ap-az2,ap-az3}  {}              {}
+ca-central-1    {ca-az1,ca-az2,ca-az3}  {}              {}
+us-east-1       {us-az1,us-az2,us-az3}  {}              {}
+
+query TT colnames
+SHOW REGIONS FROM CLUSTER
+----
+region          zones
+ap-southeast-2  {ap-az1,ap-az2,ap-az3}
+ca-central-1    {ca-az1,ca-az2,ca-az3}
+us-east-1       {us-az1,us-az2,us-az3}
+
+statement ok
+CREATE DATABASE non_multi_region_db
+
+statement ok
+CREATE DATABASE multi_region_test_db PRIMARY REGION "ca-central-1" REGIONS "ap-southeast-2", "us-east-1" SURVIVE REGION FAILURE
+
+statement ok
+USE multi_region_test_db;
+CREATE TABLE regional_primary_region_table (a int) LOCALITY REGIONAL BY TABLE IN PRIMARY REGION
+
+statement ok
+CREATE TABLE "regional_us-east-1_table" (a int) LOCALITY REGIONAL BY TABLE IN "us-east-1"
+
+statement ok
+CREATE TABLE global_table (a int) LOCALITY GLOBAL
+
+statement ok
+CREATE TABLE regional_by_row_table (
+  pk int PRIMARY KEY,
+  pk2 int NOT NULL,
+  a int NOT NULL,
+  b int NOT NULL,
+  j JSON,
+  INDEX (a),
+  UNIQUE (b),
+  INVERTED INDEX (j),
+  FAMILY (pk, pk2, a, b)
+) LOCALITY REGIONAL BY ROW
+
+statement ok
+use non_multi_region_db
+
+statement ok
+CREATE TABLE team (
+  id int PRIMARY KEY,
+  name string,
+  likes string[],
+  dislikes string[]
+)
+
+statement ok
+INSERT INTO team VALUES (1, 'arulajmani', ARRAY['turkey','coffee','ps5'], ARRAY['going outside in winter','denormalization']);
+INSERT INTO team VALUES (2, 'otan', ARRAY['Sydney suburbs','cricket','vim'], ARRAY['flaky tests','onboarding'])
+
+query ITTT colnames
+SELECT * FROM team
+----
+id  name        likes                           dislikes
+1   arulajmani  {turkey,coffee,ps5}             {"going outside in winter",denormalization}
+2   otan        {"Sydney suburbs",cricket,vim}  {"flaky tests",onboarding}
+
+statement ok
+EXPORT INTO CSV 'nodelocal://1/team_export/' WITH DELIMITER = '|' FROM TABLE team
+
+statement ok
+use multi_region_test_db;
+IMPORT TABLE team (
+  id int PRIMARY KEY,
+  name string,
+  likes string[],
+  dislikes string[]
+)
+CSV DATA ('nodelocal://1/team_export/export*.csv') WITH DELIMITER = '|'
+
+query ITTT colnames
+SELECT * FROM team
+----
+id  name        likes                           dislikes
+1   arulajmani  {turkey,coffee,ps5}             {"going outside in winter",denormalization}
+2   otan        {"Sydney suburbs",cricket,vim}  {"flaky tests",onboarding}
+
+query TT
+SHOW CREATE TABLE team
+----
+team                                            CREATE TABLE public.team (
+                                                id INT8 NOT NULL,
+                                                name STRING NULL,
+                                                likes STRING[] NULL,
+                                                dislikes STRING[] NULL,
+                                                CONSTRAINT "primary" PRIMARY KEY (id ASC),
+                                                FAMILY "primary" (id, name, likes, dislikes)
+) LOCALITY REGIONAL BY TABLE IN PRIMARY REGION