Feature/default and generated columns

evalbench/databases/spanner.py

@@ -465,7 +465,14 @@ def insert_data(self, data, setup=None):
             schema_name = 'public' if self.expected_dialect_str == "POSTGRESQL" else ''
             with self.database.snapshot() as snapshot:
                 type_col = "spanner_type" if self.expected_dialect_str == "GOOGLESQL" else "data_type"
-                query = f"SELECT table_name, column_name, {type_col} FROM information_schema.columns WHERE table_schema = '{schema_name}' ORDER BY table_name, ordinal_position"
+                query = (
+                    f"SELECT table_name, column_name, {type_col} "
+                    f"FROM information_schema.columns "
+                    f"WHERE table_schema = '{schema_name}' "
+                    f"AND (UPPER(is_generated) != 'ALWAYS' OR is_generated IS NULL) "
+                    f"AND LOWER(column_name) NOT IN ('_row_id', 'surrogate_id') "
+                    f"ORDER BY table_name, ordinal_position"
+                )
                 res = snapshot.execute_sql(query, timeout=self.query_timeout)
                 for row in res:
                     t_name, c_name, d_type = row[0], row[1], row[2]

evalbench/test/spanner_test.py

@@ -47,3 +47,63 @@ def test_get_metadata(self, client):
     def test_drop_table(self, client):
         create_table = "DROP TABLE `ut`"
         client.batch_execute([create_table])
+
+    @pytest.mark.filterwarnings("ignore::DeprecationWarning")
+    def test_generated_column(self, client):
+        # Create table with generated column
+        create_table = "CREATE TABLE `ut_gen` (id INT64, a INT64, b INT64, gen INT64 AS (a + b) STORED) PRIMARY KEY (id)"
+        client.batch_execute([create_table])
+        try:
+            # Data omitting the generated column (3 values)
+            data = {"ut_gen": [[1, 10, 20], [2, 100, 200]]}
+            client.insert_data(data)
+
+            # Verify data
+            res = client.execute("SELECT id, a, b, gen FROM `ut_gen` ORDER BY id")
+            assert len(res[0]) == 2
+            assert res[0][0]["id"] == 1
+            assert res[0][0]["gen"] == 30
+            assert res[0][1]["id"] == 2
+            assert res[0][1]["gen"] == 300
+        finally:
+            client.batch_execute(["DROP TABLE `ut_gen`"])
+
+    @pytest.mark.filterwarnings("ignore::DeprecationWarning")
+    def test_default_column_omitted(self, client):
+        # Create table with default column (using surrogate_id which is filtered out)
+        create_table = "CREATE TABLE `ut_def` (id INT64, val INT64, surrogate_id STRING(36) DEFAULT ('default-uuid')) PRIMARY KEY (id)"
+        client.batch_execute([create_table])
+        try:
+            # Data omitting default column (2 values)
+            data = {"ut_def": [[1, 10], [2, 20]]}
+            client.insert_data(data)
+
+            # Verify default value is populated
+            res = client.execute("SELECT id, val, surrogate_id FROM `ut_def` ORDER BY id")
+            assert len(res[0]) == 2
+            assert res[0][0]["id"] == 1
+            assert res[0][0]["surrogate_id"] == "default-uuid"
+            assert res[0][1]["id"] == 2
+            assert res[0][1]["surrogate_id"] == "default-uuid"
+        finally:
+            client.batch_execute(["DROP TABLE `ut_def`"])
+
+    @pytest.mark.filterwarnings("ignore::DeprecationWarning")
+    def test_default_column_included(self, client):
+        # Create table with default column
+        create_table = "CREATE TABLE `ut_def_inc` (id INT64, val INT64, def_val INT64 DEFAULT (42)) PRIMARY KEY (id)"
+        client.batch_execute([create_table])
+        try:
+            # Data including default column (3 values)
+            data = {"ut_def_inc": [[1, 10, 99], [2, 20, 100]]}
+            client.insert_data(data)
+
+            # Verify explicit value is populated
+            res = client.execute("SELECT id, val, def_val FROM `ut_def_inc` ORDER BY id")
+            assert len(res[0]) == 2
+            assert res[0][0]["id"] == 1
+            assert res[0][0]["def_val"] == 99
+            assert res[0][1]["id"] == 2
+            assert res[0][1]["def_val"] == 100
+        finally:
+            client.batch_execute(["DROP TABLE `ut_def_inc`"])