Skip to content

Commit

Permalink
chore: support named schemas (#380)
Browse files Browse the repository at this point in the history
* chore: support named schemas

Modify information_schema queries so these are able to handle databases that contain multiple schemas.

* Reformatting

---------

Co-authored-by: Ankit Agarwal <146331865+ankiaga@users.noreply.github.com>
Co-authored-by: ankiaga <ankiaga@google.com>
  • Loading branch information
3 people authored Feb 5, 2024
1 parent 62cccc3 commit 217841b
Showing 1 changed file with 30 additions and 17 deletions.
47 changes: 30 additions & 17 deletions google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py
Original file line number Diff line number Diff line change
Expand Up @@ -846,7 +846,7 @@ def get_multi_columns(
col.spanner_type, col.is_nullable, col.generation_expression
FROM information_schema.columns as col
JOIN information_schema.tables AS t
ON col.table_name = t.table_name
USING (TABLE_CATALOG, TABLE_SCHEMA, TABLE_NAME)
WHERE
{table_filter_query}
{table_type_query}
Expand Down Expand Up @@ -979,9 +979,14 @@ def get_multi_indexes(
ARRAY_AGG(ic.column_ordering)
FROM information_schema.indexes as i
JOIN information_schema.index_columns AS ic
ON ic.index_name = i.index_name AND ic.table_name = i.table_name
ON ic.index_name = i.index_name
AND ic.table_catalog = i.table_catalog
AND ic.table_schema = i.table_schema
AND ic.table_name = i.table_name
JOIN information_schema.tables AS t
ON i.table_name = t.table_name
ON i.table_catalog = t.table_catalog
AND i.table_schema = t.table_schema
AND i.table_name = t.table_name
WHERE
{table_filter_query}
{table_type_query}
Expand Down Expand Up @@ -1076,9 +1081,11 @@ def get_multi_pk_constraint(
SELECT tc.table_schema, tc.table_name, ccu.COLUMN_NAME
FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc
JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE AS ccu
ON ccu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME
USING (TABLE_CATALOG, TABLE_SCHEMA, CONSTRAINT_NAME)
JOIN information_schema.tables AS t
ON tc.table_name = t.table_name
ON tc.TABLE_CATALOG = t.TABLE_CATALOG
AND tc.TABLE_SCHEMA = t.TABLE_SCHEMA
AND tc.TABLE_NAME = t.TABLE_NAME
WHERE {table_filter_query} {table_type_query}
{schema_filter_query} tc.CONSTRAINT_TYPE = "PRIMARY KEY"
""".format(
Expand Down Expand Up @@ -1196,13 +1203,19 @@ def get_multi_foreign_keys(
)
FROM information_schema.table_constraints AS tc
JOIN information_schema.constraint_column_usage AS ccu
ON ccu.constraint_name = tc.constraint_name
USING (table_catalog, table_schema, constraint_name)
JOIN information_schema.constraint_table_usage AS ctu
ON ctu.constraint_name = tc.constraint_name
ON ctu.table_catalog = tc.table_catalog
and ctu.table_schema = tc.table_schema
and ctu.constraint_name = tc.constraint_name
JOIN information_schema.key_column_usage AS kcu
ON kcu.constraint_name = tc.constraint_name
ON kcu.table_catalog = tc.table_catalog
and kcu.table_schema = tc.table_schema
and kcu.constraint_name = tc.constraint_name
JOIN information_schema.tables AS t
ON tc.table_name = t.table_name
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}
Expand Down Expand Up @@ -1323,15 +1336,14 @@ def get_unique_constraints(self, connection, table_name, schema=None, **kw):
SELECT ccu.CONSTRAINT_NAME, ccu.COLUMN_NAME
FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc
JOIN INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE AS ccu
ON ccu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME
LEFT JOIN INFORMATION_SCHEMA.REFERENTIAL_CONSTRAINTS AS rc
on tc.CONSTRAINT_NAME = rc.CONSTRAINT_NAME
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 rc.CONSTRAINT_NAME IS NOT NULL
AND tc.CONSTRAINT_NAME IS NOT NULL
""".format(
table_name=table_name
table_schema=schema or "", table_name=table_name
)

cols = []
Expand Down Expand Up @@ -1363,10 +1375,10 @@ def has_table(self, connection, table_name, schema=None, **kw):
"""
SELECT true
FROM INFORMATION_SCHEMA.TABLES
WHERE TABLE_NAME="{table_name}"
WHERE TABLE_SCHEMA="{table_schema}" AND TABLE_NAME="{table_name}"
LIMIT 1
""".format(
table_name=table_name
table_schema=schema or "", table_name=table_name
)
)

Expand All @@ -1390,9 +1402,10 @@ def has_sequence(self, connection, sequence_name, schema=None, **kw):
SELECT true
FROM INFORMATION_SCHEMA.SEQUENCES
WHERE NAME="{sequence_name}"
AND SCHEMA="{schema}"
LIMIT 1
""".format(
sequence_name=sequence_name
sequence_name=sequence_name, schema=schema or ""
)
)

Expand Down

0 comments on commit 217841b

Please sign in to comment.