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..2864de44eea1 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,102 @@ 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, + 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 + 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" + AND tc.CONSTRAINT_NAME IS NOT NULL + GROUP BY tc.TABLE_SCHEMA, tc.TABLE_NAME, tc.CONSTRAINT_NAME + """.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]), []) + # The emulator lists the backing unique constraint of a foreign + # key once per referencing foreign key, which repeats its columns. + column_names = list(dict.fromkeys(row[3])) + table_info.append({"name": row[2], "column_names": column_names}) + + 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..9a9969786ea1 100644 --- a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py +++ b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py @@ -87,6 +87,88 @@ 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_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_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", "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 + ) + + 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 ( + "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 + + 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", "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"), []) + def test_max_size_exported(self): """Test MAX_SIZE export and int_from_size helper behavior.""" from google.cloud.sqlalchemy_spanner import MAX_SIZE