From 0b8179d3eb01542bcd5086fad9f13b58771e7300 Mon Sep 17 00:00:00 2001 From: Hui Kang Date: Sat, 24 Oct 2026 00:56:26 -0700 Subject: [PATCH 1/4] fix(sqlalchemy-spanner): group multi-column unique constraints and add get_multi_unique_constraints get_unique_constraints() returned one entry per column of a multi-column unique constraint, in no particular order. Add get_multi_unique_constraints(), which reads the unique constraints of all tables in one query with their columns in KEY_COLUMN_USAGE order, and make get_unique_constraints() use it, like the other get_*() methods. Fixes #18836 Fixes #18837 --- .../sqlalchemy_spanner/sqlalchemy_spanner.py | 97 +++++++++++++++---- .../sqlalchemy-spanner/tests/test_suite_20.py | 59 ++++++++++- .../tests/unit/test_dialect.py | 79 ++++++++++++++- 3 files changed, 211 insertions(+), 24 deletions(-) diff --git a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index 4a13d27f7b28..f85f8ab7ba5d 100644 --- a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -1669,40 +1669,101 @@ def get_table_names(self, connection, schema=None, **kw): return table_names @engine_to_connection - def get_unique_constraints(self, connection, table_name, schema=None, **kw): - """Get the table unique constraints. + def get_multi_unique_constraints( + self, connection, schema=None, filter_names=None, scope=None, kind=None, **kw + ): + """ + Return information about unique constraints in all tables + in the given schema. The method is used by SQLAlchemy introspection systems. Args: connection (sqlalchemy.engine.base.Connection): SQLAlchemy connection or engine object. - table_name (str): Name of the table to introspect. schema (str): Optional. Schema name + filter_names (Sequence[str]): Optional. Optionally return information + only for the objects listed here. + scope (sqlalchemy.engine.reflection.ObjectScope): Optional. Specifies + if constraints of default, temporary or any tables + should be reflected. Spanner does not support temporary. + kind (sqlalchemy.engine.reflection.ObjectKind): Optional. Specifies the + type of objects to reflect. Returns: - dict: Dict with the unique constraints' descriptions. + dictionary: a dictionary where the keys are two-tuple schema,table-name + and the values are list of dictionaries, each representing + a unique constraint definition. The columns of each constraint + are listed in the order of the constraint. + The schema is ``None`` if no schema is provided. """ + table_filter_query = self._get_table_filter_query(filter_names, "tc", True) + schema_filter_query = " tc.table_schema = '{schema}' AND ".format( + schema=schema or "" + ) + table_type_query = self._get_table_type_query(kind, True) + sql = """ -SELECT ccu.CONSTRAINT_NAME, ccu.COLUMN_NAME -FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc -JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE AS ccu - USING (TABLE_CATALOG, TABLE_SCHEMA, CONSTRAINT_NAME) -WHERE - tc.TABLE_NAME="{table_name}" - AND tc.TABLE_SCHEMA="{table_schema}" - AND tc.CONSTRAINT_TYPE = "UNIQUE" - AND tc.CONSTRAINT_NAME IS NOT NULL -""".format(table_schema=schema or "", table_name=table_name) + SELECT tc.table_schema, tc.table_name, tc.constraint_name, kcu.column_name + FROM information_schema.table_constraints AS tc + JOIN information_schema.key_column_usage AS kcu + ON kcu.constraint_catalog = tc.constraint_catalog + AND kcu.constraint_schema = tc.constraint_schema + AND kcu.constraint_name = tc.constraint_name + JOIN information_schema.tables AS t + ON t.table_catalog = tc.table_catalog + AND t.table_schema = tc.table_schema + AND t.table_name = tc.table_name + WHERE + {table_filter_query} + {table_type_query} + {schema_filter_query} + tc.constraint_type = "UNIQUE" + ORDER BY tc.table_catalog, tc.table_schema, tc.table_name, + tc.constraint_name, kcu.ordinal_position + """.format( + table_filter_query=table_filter_query, + table_type_query=table_type_query, + schema_filter_query=schema_filter_query, + ) - cols = [] with connection.connection.database.snapshot() as snap: - rows = snap.execute_sql(sql) + rows = list(snap.execute_sql(sql)) + result_dict = {} for row in rows: - cols.append({"name": row[0], "column_names": [row[1]]}) + table_info = result_dict.setdefault((row[0] or None, row[1]), []) + if not table_info or table_info[-1]["name"] != row[2]: + table_info.append({"name": row[2], "column_names": []}) + column_names = table_info[-1]["column_names"] + # The emulator lists the backing unique constraint of a foreign + # key once per referencing foreign key, which repeats its columns. + if row[3] not in column_names: + column_names.append(row[3]) + + return result_dict + + @engine_to_connection + def get_unique_constraints(self, connection, table_name, schema=None, **kw): + """Get the table unique constraints. - return cols + The method is used by SQLAlchemy introspection systems. + + Args: + connection (sqlalchemy.engine.base.Connection): + SQLAlchemy connection or engine object. + table_name (str): Name of the table to introspect. + schema (str): Optional. Schema name + + Returns: + list: Dicts, each of which describes a unique constraint. + """ + kind = None if not USING_SQLACLCHEMY_20 else ObjectKind.ANY + constraints = self.get_multi_unique_constraints( + connection, schema=schema, filter_names=[table_name], kind=kind + ) + schema = schema or None + return constraints.get((schema, table_name), []) @engine_to_connection def has_table(self, connection, table_name, schema=None, **kw): diff --git a/packages/sqlalchemy-spanner/tests/test_suite_20.py b/packages/sqlalchemy-spanner/tests/test_suite_20.py index 2142fef94358..a7f54d0eb2ad 100644 --- a/packages/sqlalchemy-spanner/tests/test_suite_20.py +++ b/packages/sqlalchemy-spanner/tests/test_suite_20.py @@ -1137,11 +1137,60 @@ def test_get_multi_columns( result.pop((schema, t), None) self._check_table_dict(result, exp, self._required_column_keys) - @pytest.mark.skip( - "Requires an introspection method to be implemented in SQLAlchemy first" - ) - def test_get_multi_unique_constraints(self): - pass + def test_get_multi_unique_constraints(self, metadata, connection): + """ + SPANNER OVERRIDE: + + Spanner doesn't support unique constraints, so the tables of the + original test have none. Instead, a foreign key that refers to columns + other than the primary key gets a backing unique constraint from Spanner. + Checks that it is reflected as one constraint with all its columns, also + when more than one foreign key refers to the same columns. + """ + Table( + "uq_parent", + metadata, + Column("id", sqlalchemy.Integer, primary_key=True), + Column("a", sqlalchemy.String(10)), + Column("b", sqlalchemy.String(10)), + ) + for child in ("uq_child", "uq_child2"): + Table( + child, + metadata, + Column("id", sqlalchemy.Integer, primary_key=True), + Column("x", sqlalchemy.String(10)), + Column("y", sqlalchemy.String(10)), + sqlalchemy.ForeignKeyConstraint( + ["x", "y"], ["uq_parent.a", "uq_parent.b"] + ), + ) + metadata.create_all(connection) + connection.connection.commit() + + insp = inspect(connection) + constraints = insp.get_multi_unique_constraints( + filter_names=["uq_parent", "uq_child", "uq_child2"] + ) + eq_(constraints.get((None, "uq_child"), []), []) + eq_(constraints.get((None, "uq_child2"), []), []) + eq_(len(constraints[(None, "uq_parent")]), 1) + constraint = constraints[(None, "uq_parent")][0] + eq_(constraint["column_names"], ["a", "b"]) + + insp.clear_cache() + eq_(insp.get_unique_constraints("uq_parent"), [constraint]) + eq_(insp.get_unique_constraints("uq_child"), []) + + reflected = Table("uq_parent", MetaData(), autoload_with=connection) + eq_( + [ + [col.name for col in uq.columns] + for uq in reflected.constraints + if isinstance(uq, sqlalchemy.UniqueConstraint) + ], + [["a", "b"]], + ) @pytest.mark.skip( "Requires an introspection method to be implemented in SQLAlchemy first" diff --git a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py index f44599bb4fcb..bacbac9f83e2 100644 --- a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py +++ b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py @@ -87,7 +87,84 @@ def test_get_multi_indexes_handles_null_column_orderings_array(self): index_info = res[("public", "my_table")][0] eq_(index_info["column_sorting"], {}) - def test_max_size_exported(self): + def test_get_multi_unique_constraints_groups_columns(self): + """Test get_multi_unique_constraints groups the columns of each constraint.""" + dialect = SpannerDialect() + connection = MagicMock() + mock_snapshot = MagicMock() + # Mock rows: schema, table, constraint_name, column_name, ordered by + # constraint and column position. The emulator repeats the rows of the + # backing unique constraint of a foreign key for each referencing key. + mock_snapshot.execute_sql.return_value = [ + ["", "parent", "IDX_parent_a_b_U", "a"], + ["", "parent", "IDX_parent_a_b_U", "a"], + ["", "parent", "IDX_parent_a_b_U", "b"], + ["", "parent", "IDX_parent_a_b_U", "b"], + ["", "parent", "IDX_parent_c_U", "c"], + ["my_schema", "other", "IDX_other_x_U", "x"], + ] + connection.connection.database.snapshot.return_value.__enter__.return_value = ( + mock_snapshot + ) + + res = dialect.get_multi_unique_constraints(connection) + + eq_( + res, + { + (None, "parent"): [ + {"name": "IDX_parent_a_b_U", "column_names": ["a", "b"]}, + {"name": "IDX_parent_c_U", "column_names": ["c"]}, + ], + ("my_schema", "other"): [ + {"name": "IDX_other_x_U", "column_names": ["x"]} + ], + }, + ) + + def test_get_multi_unique_constraints_sql(self): + """Test get_multi_unique_constraints reads all tables in one query.""" + dialect = SpannerDialect() + connection = MagicMock() + mock_snapshot = MagicMock() + mock_snapshot.execute_sql.return_value = [] + connection.connection.database.snapshot.return_value.__enter__.return_value = ( + mock_snapshot + ) + + dialect.get_multi_unique_constraints( + connection, filter_names=["parent", "other"] + ) + + mock_snapshot.execute_sql.assert_called_once() + executed_sql = mock_snapshot.execute_sql.call_args[0][0] + assert "information_schema.key_column_usage AS kcu" in executed_sql + assert "kcu.ordinal_position" in executed_sql + assert "tc.table_name = 'parent' OR tc.table_name = 'other'" in executed_sql + assert 'tc.constraint_type = "UNIQUE"' in executed_sql + + def test_get_unique_constraints(self): + """Test get_unique_constraints returns the constraints of one table.""" + dialect = SpannerDialect() + connection = MagicMock() + mock_snapshot = MagicMock() + mock_snapshot.execute_sql.return_value = [ + ["", "parent", "IDX_parent_a_b_U", "a"], + ["", "parent", "IDX_parent_a_b_U", "b"], + ] + connection.connection.database.snapshot.return_value.__enter__.return_value = ( + mock_snapshot + ) + + eq_( + dialect.get_unique_constraints(connection, "parent"), + [{"name": "IDX_parent_a_b_U", "column_names": ["a", "b"]}], + ) + executed_sql = mock_snapshot.execute_sql.call_args[0][0] + assert "tc.table_name = 'parent'" in executed_sql + + mock_snapshot.execute_sql.return_value = [] + eq_(dialect.get_unique_constraints(connection, "no_constraints"), []) """Test MAX_SIZE export and int_from_size helper behavior.""" from google.cloud.sqlalchemy_spanner import MAX_SIZE from google.cloud.sqlalchemy_spanner.sqlalchemy_spanner import ( From efbdf2975631e0773644b8dad5a745c89220fe00 Mon Sep 17 00:00:00 2001 From: Hui Kang Date: Sat, 24 Oct 2026 02:07:54 -0700 Subject: [PATCH 2/4] test(sqlalchemy-spanner): restore test_max_size_exported The previous commit dropped its def line, so its body ran as part of test_get_unique_constraints. --- packages/sqlalchemy-spanner/tests/unit/test_dialect.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py index bacbac9f83e2..e54e0c179e9e 100644 --- a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py +++ b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py @@ -165,6 +165,8 @@ def test_get_unique_constraints(self): mock_snapshot.execute_sql.return_value = [] eq_(dialect.get_unique_constraints(connection, "no_constraints"), []) + + def test_max_size_exported(self): """Test MAX_SIZE export and int_from_size helper behavior.""" from google.cloud.sqlalchemy_spanner import MAX_SIZE from google.cloud.sqlalchemy_spanner.sqlalchemy_spanner import ( From 8ff885b6aee9f04ab0cf6c02d477e852c59f5ca0 Mon Sep 17 00:00:00 2001 From: Hui Kang Date: Sat, 24 Oct 2026 02:33:59 -0700 Subject: [PATCH 3/4] refactor(sqlalchemy-spanner): keep the CONSTRAINT_NAME filter in the unique constraints query Also write the INFORMATION_SCHEMA table and column names in upper case, as the replaced query did. --- .../sqlalchemy_spanner/sqlalchemy_spanner.py | 29 ++++++++++--------- .../tests/unit/test_dialect.py | 7 +++-- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index f85f8ab7ba5d..81517da3ac12 100644 --- a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -1698,29 +1698,30 @@ def get_multi_unique_constraints( The schema is ``None`` if no schema is provided. """ table_filter_query = self._get_table_filter_query(filter_names, "tc", True) - schema_filter_query = " tc.table_schema = '{schema}' AND ".format( + schema_filter_query = " tc.TABLE_SCHEMA = '{schema}' AND ".format( schema=schema or "" ) table_type_query = self._get_table_type_query(kind, True) sql = """ - SELECT tc.table_schema, tc.table_name, tc.constraint_name, kcu.column_name - FROM information_schema.table_constraints AS tc - JOIN information_schema.key_column_usage AS kcu - ON kcu.constraint_catalog = tc.constraint_catalog - AND kcu.constraint_schema = tc.constraint_schema - AND kcu.constraint_name = tc.constraint_name - JOIN information_schema.tables AS t - ON t.table_catalog = tc.table_catalog - AND t.table_schema = tc.table_schema - AND t.table_name = tc.table_name + SELECT tc.TABLE_SCHEMA, tc.TABLE_NAME, tc.CONSTRAINT_NAME, kcu.COLUMN_NAME + FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc + JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS kcu + ON kcu.CONSTRAINT_CATALOG = tc.CONSTRAINT_CATALOG + AND kcu.CONSTRAINT_SCHEMA = tc.CONSTRAINT_SCHEMA + AND kcu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME + JOIN INFORMATION_SCHEMA.TABLES AS t + ON t.TABLE_CATALOG = tc.TABLE_CATALOG + AND t.TABLE_SCHEMA = tc.TABLE_SCHEMA + AND t.TABLE_NAME = tc.TABLE_NAME WHERE {table_filter_query} {table_type_query} {schema_filter_query} - tc.constraint_type = "UNIQUE" - ORDER BY tc.table_catalog, tc.table_schema, tc.table_name, - tc.constraint_name, kcu.ordinal_position + tc.CONSTRAINT_TYPE = "UNIQUE" + AND tc.CONSTRAINT_NAME IS NOT NULL + ORDER BY tc.TABLE_CATALOG, tc.TABLE_SCHEMA, tc.TABLE_NAME, + tc.CONSTRAINT_NAME, kcu.ORDINAL_POSITION """.format( table_filter_query=table_filter_query, table_type_query=table_type_query, diff --git a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py index e54e0c179e9e..5e8f0f287682 100644 --- a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py +++ b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py @@ -138,10 +138,11 @@ def test_get_multi_unique_constraints_sql(self): mock_snapshot.execute_sql.assert_called_once() executed_sql = mock_snapshot.execute_sql.call_args[0][0] - assert "information_schema.key_column_usage AS kcu" in executed_sql - assert "kcu.ordinal_position" in executed_sql + assert "INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS kcu" in executed_sql + assert "kcu.ORDINAL_POSITION" in executed_sql assert "tc.table_name = 'parent' OR tc.table_name = 'other'" in executed_sql - assert 'tc.constraint_type = "UNIQUE"' in executed_sql + assert 'tc.CONSTRAINT_TYPE = "UNIQUE"' in executed_sql + assert "tc.CONSTRAINT_NAME IS NOT NULL" in executed_sql def test_get_unique_constraints(self): """Test get_unique_constraints returns the constraints of one table.""" From 3a56358bf31e41f1850d8c54685fd2c77ff504b5 Mon Sep 17 00:00:00 2001 From: Hui Kang Date: Sat, 24 Oct 2026 02:53:21 -0700 Subject: [PATCH 4/4] refactor(sqlalchemy-spanner): aggregate unique constraint columns in SQL Group by schema, table and constraint name and collect the columns with ARRAY_AGG(... ORDER BY ORDINAL_POSITION), instead of ordering one row per column and grouping them in Python. --- .../sqlalchemy_spanner/sqlalchemy_spanner.py | 16 ++++++------ .../tests/unit/test_dialect.py | 26 ++++++++++--------- 2 files changed, 22 insertions(+), 20 deletions(-) diff --git a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index 81517da3ac12..2864de44eea1 100644 --- a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -1704,7 +1704,11 @@ def get_multi_unique_constraints( table_type_query = self._get_table_type_query(kind, True) sql = """ - SELECT tc.TABLE_SCHEMA, tc.TABLE_NAME, tc.CONSTRAINT_NAME, kcu.COLUMN_NAME + SELECT + tc.TABLE_SCHEMA, + tc.TABLE_NAME, + tc.CONSTRAINT_NAME, + ARRAY_AGG(kcu.COLUMN_NAME ORDER BY kcu.ORDINAL_POSITION) FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS kcu ON kcu.CONSTRAINT_CATALOG = tc.CONSTRAINT_CATALOG @@ -1720,8 +1724,7 @@ def get_multi_unique_constraints( {schema_filter_query} tc.CONSTRAINT_TYPE = "UNIQUE" AND tc.CONSTRAINT_NAME IS NOT NULL - ORDER BY tc.TABLE_CATALOG, tc.TABLE_SCHEMA, tc.TABLE_NAME, - tc.CONSTRAINT_NAME, kcu.ORDINAL_POSITION + GROUP BY tc.TABLE_SCHEMA, tc.TABLE_NAME, tc.CONSTRAINT_NAME """.format( table_filter_query=table_filter_query, table_type_query=table_type_query, @@ -1734,13 +1737,10 @@ def get_multi_unique_constraints( for row in rows: table_info = result_dict.setdefault((row[0] or None, row[1]), []) - if not table_info or table_info[-1]["name"] != row[2]: - table_info.append({"name": row[2], "column_names": []}) - column_names = table_info[-1]["column_names"] # The emulator lists the backing unique constraint of a foreign # key once per referencing foreign key, which repeats its columns. - if row[3] not in column_names: - column_names.append(row[3]) + column_names = list(dict.fromkeys(row[3])) + table_info.append({"name": row[2], "column_names": column_names}) return result_dict diff --git a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py index 5e8f0f287682..9a9969786ea1 100644 --- a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py +++ b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py @@ -92,16 +92,13 @@ def test_get_multi_unique_constraints_groups_columns(self): dialect = SpannerDialect() connection = MagicMock() mock_snapshot = MagicMock() - # Mock rows: schema, table, constraint_name, column_name, ordered by - # constraint and column position. The emulator repeats the rows of the - # backing unique constraint of a foreign key for each referencing key. + # Mock rows: schema, table, constraint_name, column_names. The emulator + # repeats the columns of the backing unique constraint of a foreign key + # for each referencing foreign key. mock_snapshot.execute_sql.return_value = [ - ["", "parent", "IDX_parent_a_b_U", "a"], - ["", "parent", "IDX_parent_a_b_U", "a"], - ["", "parent", "IDX_parent_a_b_U", "b"], - ["", "parent", "IDX_parent_a_b_U", "b"], - ["", "parent", "IDX_parent_c_U", "c"], - ["my_schema", "other", "IDX_other_x_U", "x"], + ["", "parent", "IDX_parent_a_b_U", ["a", "a", "b", "b"]], + ["", "parent", "IDX_parent_c_U", ["c"]], + ["my_schema", "other", "IDX_other_x_U", ["x"]], ] connection.connection.database.snapshot.return_value.__enter__.return_value = ( mock_snapshot @@ -139,7 +136,13 @@ def test_get_multi_unique_constraints_sql(self): mock_snapshot.execute_sql.assert_called_once() executed_sql = mock_snapshot.execute_sql.call_args[0][0] assert "INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS kcu" in executed_sql - assert "kcu.ORDINAL_POSITION" in executed_sql + assert ( + "ARRAY_AGG(kcu.COLUMN_NAME ORDER BY kcu.ORDINAL_POSITION)" in executed_sql + ) + assert ( + "GROUP BY tc.TABLE_SCHEMA, tc.TABLE_NAME, tc.CONSTRAINT_NAME" + in executed_sql + ) assert "tc.table_name = 'parent' OR tc.table_name = 'other'" in executed_sql assert 'tc.CONSTRAINT_TYPE = "UNIQUE"' in executed_sql assert "tc.CONSTRAINT_NAME IS NOT NULL" in executed_sql @@ -150,8 +153,7 @@ def test_get_unique_constraints(self): connection = MagicMock() mock_snapshot = MagicMock() mock_snapshot.execute_sql.return_value = [ - ["", "parent", "IDX_parent_a_b_U", "a"], - ["", "parent", "IDX_parent_a_b_U", "b"], + ["", "parent", "IDX_parent_a_b_U", ["a", "b"]], ] connection.connection.database.snapshot.return_value.__enter__.return_value = ( mock_snapshot