diff --git a/evalbench/databases/spanner.py b/evalbench/databases/spanner.py index 24c4628d..60011033 100644 --- a/evalbench/databases/spanner.py +++ b/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] diff --git a/evalbench/test/spanner_test.py b/evalbench/test/spanner_test.py index 200122b0..65fd4ebc 100644 --- a/evalbench/test/spanner_test.py +++ b/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`"])