[tests] Add integration tests for IDENTITY columns to bulk, live, and reverse Spanner migration templates (#2992)

nmemond · web-flow · commit 793fa0fd15e5 · 2025-11-18T12:46:57.000+05:30
* [tests] Add integration tests for IDENTITY columns to bulk, live, and reverse Spanner migration templates

* Actually add missing test resource...
diff --git a/v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamToSpannerDDLIT.java b/v2/datastream-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamToSpannerDDLIT.java
@@ -463,6 +463,36 @@ public void migrationTestWithSequenceColumns() {
     assertThatResult(result).meetsConditions();
   }
 
+  @Test
+  public void migrationTestWithIdentityColumns() {
+    // Construct a ChainedConditionCheck with 2 stages.
+    // 1. Send initial wave of events
+    // 2. Wait on Spanner to have events
+    ChainedConditionCheck conditionCheck =
+        ChainedConditionCheck.builder(
+                List.of(
+                    uploadDataStreamFile(
+                        jobInfo,
+                        TABLE6,
+                        "identity.avro",
+                        "DataStreamToSpannerDDLIT/Books.avro",
+                        gcsResourceManager),
+                    SpannerRowsCheck.builder(spannerResourceManager, TABLE6)
+                        .setMinRows(3)
+                        .setMaxRows(3)
+                        .build()))
+            .build();
+
+    // Wait for conditions
+    PipelineOperator.Result result =
+        pipelineOperator()
+            .waitForCondition(createConfig(jobInfo, Duration.ofMinutes(8)), conditionCheck);
+
+    // Assert Conditions
+    assertThatResult(result).meetsConditions();
+    assertBooksBackfillContents();
+  }
+
   private void assertAllDatatypeColumnsTableBackfillContents() {
     List<Map<String, Object>> events = new ArrayList<>();
 
@@ -951,4 +981,26 @@ private void assertAuthorsBackfillContents() {
             spannerResourceManager.runQuery("select id, name from Authors"))
         .hasRecordsUnorderedCaseInsensitiveColumns(events);
   }
+
+  private void assertBooksBackfillContents() {
+    List<Map<String, Object>> events = new ArrayList<>();
+
+    Map<String, Object> row = new HashMap<>();
+    row.put("id", 1);
+    row.put("title", "The Lord of the Rings");
+    events.add(row);
+
+    row = new HashMap<>();
+    row.put("id", 2);
+    row.put("title", "Pride and Prejudice");
+    events.add(row);
+
+    row = new HashMap<>();
+    row.put("id", 3);
+    row.put("title", "The Hitchhiker's Guide to the Galaxy");
+    events.add(row);
+
+    SpannerAsserts.assertThatStructs(spannerResourceManager.runQuery("select id, title from Books"))
+        .hasRecordsUnorderedCaseInsensitiveColumns(events);
+  }
 }
diff --git a/v2/datastream-to-spanner/src/test/resources/DataStreamToSpannerDDLIT/Books.avro b/v2/datastream-to-spanner/src/test/resources/DataStreamToSpannerDDLIT/Books.avro
diff --git a/v2/datastream-to-spanner/src/test/resources/DataStreamToSpannerDDLIT/mysql-schema.sql b/v2/datastream-to-spanner/src/test/resources/DataStreamToSpannerDDLIT/mysql-schema.sql
@@ -129,4 +129,10 @@ CREATE TABLE `Singers` (
     `first_name` varchar(1024),
     `last_name` varchar(1024),
     PRIMARY KEY (`singer_id`)
-);
+);
+
+CREATE TABLE `Books` (
+    `id` int NOT NULL AUTO_INCREMENT,
+    `title` varchar(200),
+    PRIMARY KEY (`id`)
+) DEFAULT CHARSET=latin1;
diff --git a/v2/datastream-to-spanner/src/test/resources/DataStreamToSpannerDDLIT/spanner-schema.sql b/v2/datastream-to-spanner/src/test/resources/DataStreamToSpannerDDLIT/spanner-schema.sql
@@ -124,3 +124,8 @@ CREATE TABLE IF NOT EXISTS Singers (
     first_name STRING(1024),
     last_name STRING(1024),
 ) PRIMARY KEY (singer_id);
+
+CREATE TABLE IF NOT EXISTS Books (
+   id INT64 NOT NULL GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
+   title STRING(200),
+) PRIMARY KEY (id);
diff --git a/v2/sourcedb-to-spanner/src/test/resources/DDLIT/company-mysql-schema.sql b/v2/sourcedb-to-spanner/src/test/resources/DDLIT/company-mysql-schema.sql
@@ -10,7 +10,7 @@ INSERT INTO `company` VALUES
                           (3,'ama','1994-07-05');
 
 CREATE TABLE `employee` (
-                            `employee_id` int(11) PRIMARY KEY NOT NULL,
+                            `employee_id` int(11) AUTO_INCREMENT PRIMARY KEY NOT NULL,
                             `company_id` int(11) DEFAULT NULL,
                             `employee_name` varchar(100) DEFAULT NULL,
                             `employee_address` varchar(100) DEFAULT NULL,
diff --git a/v2/sourcedb-to-spanner/src/test/resources/DDLIT/company-session.json b/v2/sourcedb-to-spanner/src/test/resources/DDLIT/company-session.json
@@ -97,8 +97,8 @@
           "Comment": "From: employee_id int(10)",
           "Id": "c11",
           "AutoGen": {
-            "Name": "",
-            "GenerationType": ""
+            "Name": "Identity",
+            "GenerationType": "Identity"
           }
         },
         "c12": {
@@ -507,7 +507,11 @@
             "ForeignKey": false,
             "AutoIncrement": false
           },
-          "Id": "c11"
+          "Id": "c11",
+          "AutoGen": {
+            "Name": "Auto Increment",
+            "GenerationType": "Auto Increment"
+          }
         },
         "c12": {
           "Name": "company_id",
diff --git a/v2/sourcedb-to-spanner/src/test/resources/DDLIT/company-spanner-schema.sql b/v2/sourcedb-to-spanner/src/test/resources/DDLIT/company-spanner-schema.sql
@@ -9,7 +9,7 @@ CREATE TABLE IF NOT EXISTS
 CREATE TABLE IF NOT EXISTS
     employee
 (
-    employee_id         INT64 NOT NULL,
+    employee_id         INT64 NOT NULL GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
     company_id          INT64,
     employee_name       STRING(100),
     employee_address    STRING(100),
diff --git a/v2/spanner-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/SpannerToSourceDbIT.java b/v2/spanner-to-sourcedb/src/test/java/com/google/cloud/teleport/v2/templates/SpannerToSourceDbIT.java
@@ -75,6 +75,7 @@ public class SpannerToSourceDbIT extends SpannerToSourceDbITBase {
   private static final String TABLE = "Users";
   private static final String TABLE_WITH_VIRTUAL_GEN_COL = "TableWithVirtualGeneratedColumn";
   private static final String TABLE_WITH_STORED_GEN_COL = "TableWithStoredGeneratedColumn";
+  private static final String TABLE_WITH_IDENTITY_COL = "TableWithIdentityColumn";
   private static final String BOUNDARY_CHECK_TABLE =
       "testtable_03TpCoVF16ED0KLxM3v808cH3bTGQ0uK_FEXuZHbttvYZPAeGeqiO";
   private static final HashSet<SpannerToSourceDbIT> testInstances = new HashSet<>();
@@ -273,6 +274,23 @@ public void spannerToMySQLSourceDbMaxColAndTableNameTest()
     assertBoundaryRowInMySQL();
   }
 
+  @Test
+  public void spannerToSourceDbWithIdentityColumns() {
+    assertThatPipeline(jobInfo).isRunning();
+    // INSERT
+    writeRowsWithIdentityColInSpanner();
+
+    assertThatPipeline(jobInfo).isRunning();
+
+    PipelineOperator.Result result =
+        pipelineOperator()
+            .waitForCondition(
+                createConfig(jobInfo, TEST_TIMEOUT),
+                () -> jdbcResourceManager.getRowCount(TABLE_WITH_IDENTITY_COL) == 2);
+    assertThatResult(result).meetsConditions();
+    assertIdentityColRowsInMySQLAfterInsert();
+  }
+
   private void writeMaxColRowsInSpanner() {
     List<Mutation> mutations = new ArrayList<>();
     Mutation.WriteBuilder mutationBuilder =
@@ -413,4 +431,33 @@ private boolean allGenColRowsDeleted() {
     long rowCountTable2 = jdbcResourceManager.getRowCount(TABLE_WITH_VIRTUAL_GEN_COL);
     return (rowCountTable1 == 0) && (rowCountTable2 == 0);
   }
+
+  private void writeRowsWithIdentityColInSpanner() {
+    List<Mutation> mutations = new ArrayList<>();
+    mutations.add(
+        Mutation.newInsertBuilder(TABLE_WITH_IDENTITY_COL)
+            .set("id")
+            .to(1)
+            .set("column1")
+            .to("id1")
+            .build());
+    mutations.add(
+        Mutation.newInsertBuilder(TABLE_WITH_IDENTITY_COL)
+            .set("id")
+            .to(2)
+            .set("column1")
+            .to("id2")
+            .build());
+
+    spannerResourceManager.write(mutations);
+  }
+
+  private void assertIdentityColRowsInMySQLAfterInsert() {
+    List<Map<String, Object>> rows = jdbcResourceManager.readTable(TABLE_WITH_IDENTITY_COL);
+    assertThat(rows).hasSize(2);
+    assertThat(rows.get(0).get("id")).isEqualTo(1);
+    assertThat(rows.get(0).get("column1")).isEqualTo("id1");
+    assertThat(rows.get(1).get("id")).isEqualTo(2);
+    assertThat(rows.get(1).get("column1")).isEqualTo("id2");
+  }
 }
diff --git a/v2/spanner-to-sourcedb/src/test/resources/SpannerToSourceDbIT/mysql-schema.sql b/v2/spanner-to-sourcedb/src/test/resources/SpannerToSourceDbIT/mysql-schema.sql
@@ -22,3 +22,9 @@ CREATE TABLE Users (
  id INT NOT NULL,
  col_qcbF69RmXTRe3B_03TpCoVF16ED0KLxM3v808cH3bTGQ0uK_FEXuZHbttvY VARCHAR(25),
  PRIMARY KEY(id));
+
+ CREATE TABLE TableWithIdentityColumn (
+     id BIGINT NOT NULL AUTO_INCREMENT,
+     column1 VARCHAR(25),
+     PRIMARY KEY(id)
+ );
diff --git a/v2/spanner-to-sourcedb/src/test/resources/SpannerToSourceDbIT/session.json b/v2/spanner-to-sourcedb/src/test/resources/SpannerToSourceDbIT/session.json
@@ -363,6 +363,79 @@
       "CheckConstraints": null,
       "Comment": "Spanner schema for source table TableWithVirtualGeneratedColumn",
       "Id": "t47"
+    },
+    "t48": {
+      "Name": "TableWithIdentityColumn",
+      "ColIds": [
+        "c54",
+        "c55"
+      ],
+      "ShardIdColumn": "",
+      "ColDefs": {
+        "c54": {
+          "Name": "id",
+          "T": {
+            "Name": "INT64",
+            "Len": 0,
+            "IsArray": false
+          },
+          "NotNull": true,
+          "Comment": "From: id int(10)",
+          "Id": "c54",
+          "AutoGen": {
+            "Name": "Identity",
+            "GenerationType": "Identity"
+          },
+          "DefaultValue": {
+            "IsPresent": false,
+            "Value": {
+              "ExpressionId": "",
+              "Statement": ""
+            }
+          },
+          "Opts": null
+        },
+        "c55": {
+          "Name": "column1",
+          "T": {
+            "Name": "STRING",
+            "Len": 25,
+            "IsArray": false
+          },
+          "NotNull": false,
+          "Comment": "From: column1 varchar(25)",
+          "Id": "c55",
+          "AutoGen": {
+            "Name": "",
+            "GenerationType": ""
+          },
+          "DefaultValue": {
+            "IsPresent": false,
+            "Value": {
+              "ExpressionId": "",
+              "Statement": ""
+            }
+          },
+          "Opts": null
+        }
+      },
+      "PrimaryKeys": [
+        {
+          "ColId": "c54",
+          "Desc": false,
+          "Order": 1
+        }
+      ],
+      "ForeignKeys": null,
+      "Indexes": null,
+      "ParentTable": {
+        "Id": "",
+        "OnDelete": "",
+        "InterleaveType": ""
+      },
+      "CheckConstraints": null,
+      "Comment": "Spanner schema for source table TableWithIdentityColumn",
+      "Id": "t48"
     }
   },
   "SyntheticPKeys": {},
@@ -794,6 +867,89 @@
       "CheckConstraints": null,
       "Indexes": null,
       "Id": "t47"
+    },
+    "t48": {
+      "Name": "TableWithIdentityColumn",
+      "Schema": "test_limits",
+      "ColIds": [
+        "c54",
+        "c55"
+      ],
+      "ColDefs": {
+        "c54": {
+          "Name": "id",
+          "Type": {
+            "Name": "int",
+            "Mods": [
+              10
+            ],
+            "ArrayBounds": null
+          },
+          "NotNull": true,
+          "Ignored": {
+            "Check": false,
+            "Identity": false,
+            "Default": false,
+            "Exclusion": false,
+            "ForeignKey": false,
+            "AutoIncrement": false
+          },
+          "Id": "c54",
+          "AutoGen": {
+            "Name": "Auto Increment",
+            "GenerationType": "Auto Increment"
+          },
+          "DefaultValue": {
+            "IsPresent": false,
+            "Value": {
+              "ExpressionId": "",
+              "Statement": ""
+            }
+          }
+        },
+        "c55": {
+          "Name": "column1",
+          "Type": {
+            "Name": "varchar",
+            "Mods": [
+              25
+            ],
+            "ArrayBounds": null
+          },
+          "NotNull": false,
+          "Ignored": {
+            "Check": false,
+            "Identity": false,
+            "Default": false,
+            "Exclusion": false,
+            "ForeignKey": false,
+            "AutoIncrement": false
+          },
+          "Id": "c55",
+          "AutoGen": {
+            "Name": "",
+            "GenerationType": ""
+          },
+          "DefaultValue": {
+            "IsPresent": false,
+            "Value": {
+              "ExpressionId": "",
+              "Statement": ""
+            }
+          }
+        }
+      },
+      "PrimaryKeys": [
+        {
+          "ColId": "c54",
+          "Desc": false,
+          "Order": 1
+        }
+      ],
+      "ForeignKeys": null,
+      "CheckConstraints": null,
+      "Indexes": null,
+      "Id": "t48"
     }
   },
   "SchemaIssues": {
@@ -838,6 +994,10 @@
         ]
       },
       "TableLevelIssues": null
+    },
+    "t48": {
+      "ColumnLevelIssues": {},
+      "TableLevelIssues": null
     }
   },
   "InvalidCheckExp": null,
diff --git a/v2/spanner-to-sourcedb/src/test/resources/SpannerToSourceDbIT/spanner-schema.sql b/v2/spanner-to-sourcedb/src/test/resources/SpannerToSourceDbIT/spanner-schema.sql
@@ -26,6 +26,11 @@ CREATE TABLE TableWithStoredGeneratedColumn (
  col_qcbF69RmXTRe3B_03TpCoVF16ED0KLxM3v808cH3bTGQ0uK_FEXuZHbttvY STRING(25),
  ) PRIMARY KEY(id);
 
+CREATE TABLE TableWithIdentityColumn (
+    id INT64 NOT NULL GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),
+    column1 STRING(25),
+) PRIMARY KEY(id);
+
 CREATE CHANGE STREAM allstream
   FOR ALL OPTIONS (
   value_capture_type = 'NEW_ROW',

Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,7 @@ CREATE TABLE IF NOT EXISTS`
`9`	`9`	`CREATE TABLE IF NOT EXISTS`
`10`	`10`	`employee`
`11`	`11`	`(`
`12`		`- employee_id INT64 NOT NULL,`
	`12`	`+ employee_id INT64 NOT NULL GENERATED BY DEFAULT AS IDENTITY (BIT_REVERSED_POSITIVE),`
`13`	`13`	`company_id INT64,`
`14`	`14`	`employee_name STRING(100),`
`15`	`15`	`employee_address STRING(100),`