Skip to content

Bracket-quote the database name in SQL Server catalog queries #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
maharanay22 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Bracket-quote the database name in SQL Server catalog queries #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 28 additions & 23 deletions mindsdb/integrations/handlers/mssql_handler/mssql_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,11 @@ def __del__(self):
if self.is_connected is True:
self.disconnect()

@property
def _quoted_database(self) -> str:
"""The database name as a bracket-quoted T-SQL identifier, so names like 'sqldb-prod' work in queries."""
return "[" + str(self.database).replace("]", "]]") + "]"

def connect(self):
"""
Establishes a connection to a Microsoft SQL Server database.
Expand Down Expand Up @@ -448,7 +453,7 @@ def get_tables(self) -> Response:
table_schema,
table_name,
table_type
FROM {self.database}.INFORMATION_SCHEMA.TABLES
FROM {self._quoted_database}.INFORMATION_SCHEMA.TABLES
WHERE TABLE_TYPE in ('BASE TABLE', 'VIEW')
"""
if self.schema:
Expand Down Expand Up @@ -514,17 +519,17 @@ def meta_get_tables(self, table_names: list[str] | None = None) -> Response:
t.TABLE_TYPE as table_type,
CAST(ep.value AS NVARCHAR(MAX)) as table_description,
SUM(p.rows) as row_count
FROM {self.database}.INFORMATION_SCHEMA.TABLES t
LEFT JOIN {self.database}.sys.tables st
FROM {self._quoted_database}.INFORMATION_SCHEMA.TABLES t
LEFT JOIN {self._quoted_database}.sys.tables st
ON t.TABLE_NAME = st.name
LEFT JOIN {self.database}.sys.schemas s
LEFT JOIN {self._quoted_database}.sys.schemas s
ON st.schema_id = s.schema_id AND t.TABLE_SCHEMA = s.name
LEFT JOIN {self.database}.sys.extended_properties ep
LEFT JOIN {self._quoted_database}.sys.extended_properties ep
ON st.object_id = ep.major_id
AND ep.minor_id = 0
AND ep.class = 1
AND ep.name = 'MS_Description'
LEFT JOIN {self.database}.sys.partitions p
LEFT JOIN {self._quoted_database}.sys.partitions p
ON st.object_id = p.object_id
AND p.index_id IN (0, 1)
WHERE t.TABLE_TYPE IN ('BASE TABLE', 'VIEW')
Expand Down Expand Up @@ -561,14 +566,14 @@ def meta_get_columns(self, table_names: list[str] | None = None) -> Response:
CAST(ep.value AS NVARCHAR(MAX)) as column_description,
c.COLUMN_DEFAULT as column_default,
CASE WHEN c.IS_NULLABLE = 'YES' THEN 1 ELSE 0 END as is_nullable
FROM {self.database}.INFORMATION_SCHEMA.COLUMNS c
LEFT JOIN {self.database}.sys.tables st
FROM {self._quoted_database}.INFORMATION_SCHEMA.COLUMNS c
LEFT JOIN {self._quoted_database}.sys.tables st
ON c.TABLE_NAME = st.name
LEFT JOIN {self.database}.sys.schemas s
LEFT JOIN {self._quoted_database}.sys.schemas s
ON st.schema_id = s.schema_id AND c.TABLE_SCHEMA = s.name
LEFT JOIN {self.database}.sys.columns sc
LEFT JOIN {self._quoted_database}.sys.columns sc
ON st.object_id = sc.object_id AND c.COLUMN_NAME = sc.name
LEFT JOIN {self.database}.sys.extended_properties ep
LEFT JOIN {self._quoted_database}.sys.extended_properties ep
ON st.object_id = ep.major_id
AND sc.column_id = ep.minor_id
AND ep.name = 'MS_Description'
Expand Down Expand Up @@ -625,14 +630,14 @@ def meta_get_column_statistics(self, table_names: list[str] | None = None) -> Re
NULL AS MOST_COMMON_FREQUENCIES,
CAST(h.min_value AS NVARCHAR(MAX)) AS MINIMUM_VALUE,
CAST(h.max_value AS NVARCHAR(MAX)) AS MAXIMUM_VALUE
FROM {self.database}.sys.tables t
INNER JOIN {self.database}.sys.schemas s
FROM {self._quoted_database}.sys.tables t
INNER JOIN {self._quoted_database}.sys.schemas s
ON t.schema_id = s.schema_id
INNER JOIN {self.database}.sys.columns c
INNER JOIN {self._quoted_database}.sys.columns c
ON t.object_id = c.object_id
LEFT JOIN {self.database}.sys.stats st
LEFT JOIN {self._quoted_database}.sys.stats st
ON st.object_id = t.object_id
LEFT JOIN {self.database}.sys.stats_columns sc
LEFT JOIN {self._quoted_database}.sys.stats_columns sc
ON sc.object_id = st.object_id
AND sc.stats_id = st.stats_id
AND sc.column_id = c.column_id
Expand All @@ -642,7 +647,7 @@ def meta_get_column_statistics(self, table_names: list[str] | None = None) -> Re
MIN(CAST(range_high_key AS NVARCHAR(MAX))) AS min_value,
MAX(CAST(range_high_key AS NVARCHAR(MAX))) AS max_value,
SUM(CAST(distinct_range_rows AS BIGINT)) + COUNT(*) AS distinct_count
FROM {self.database}.sys.dm_db_stats_histogram(st.object_id, st.stats_id)
FROM {self._quoted_database}.sys.dm_db_stats_histogram(st.object_id, st.stats_id)
WHERE st.object_id IS NOT NULL
) h
WHERE s.name NOT IN ('sys', 'INFORMATION_SCHEMA')
Expand Down Expand Up @@ -670,8 +675,8 @@ def meta_get_primary_keys(self, table_names: list[str] | None = None) -> Respons
kcu.COLUMN_NAME as column_name,
kcu.ORDINAL_POSITION as ordinal_position,
tc.CONSTRAINT_NAME as constraint_name
FROM {self.database}.INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc
INNER JOIN {self.database}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE kcu
FROM {self._quoted_database}.INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc
INNER JOIN {self._quoted_database}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE kcu
ON tc.CONSTRAINT_NAME = kcu.CONSTRAINT_NAME
AND tc.TABLE_SCHEMA = kcu.TABLE_SCHEMA
AND tc.TABLE_NAME = kcu.TABLE_NAME
Expand Down Expand Up @@ -707,12 +712,12 @@ def meta_get_foreign_keys(self, table_names: list[str] | None = None) -> Respons
OBJECT_NAME(fk.parent_object_id) as child_table_name,
COL_NAME(fkc.parent_object_id, fkc.parent_column_id) as child_column_name,
fk.name as constraint_name
FROM {self.database}.sys.foreign_keys fk
INNER JOIN {self.database}.sys.foreign_key_columns fkc
FROM {self._quoted_database}.sys.foreign_keys fk
INNER JOIN {self._quoted_database}.sys.foreign_key_columns fkc
ON fk.object_id = fkc.constraint_object_id
INNER JOIN {self.database}.sys.tables t
INNER JOIN {self._quoted_database}.sys.tables t
ON fk.parent_object_id = t.object_id
INNER JOIN {self.database}.sys.schemas s
INNER JOIN {self._quoted_database}.sys.schemas s
ON t.schema_id = s.schema_id
WHERE s.name NOT IN ('sys', 'INFORMATION_SCHEMA')
"""
Expand Down
24 changes: 23 additions & 1 deletion tests/unit/handlers/test_mssql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,12 +193,34 @@ def test_get_tables(self):
call_args = self.handler.native_query.call_args[0][0]
database = self.handler.connection_args["database"]

self.assertIn(f"{database}.INFORMATION_SCHEMA.TABLES", call_args)
self.assertIn(f"[{database}].INFORMATION_SCHEMA.TABLES", call_args)
self.assertIn("table_schema", call_args)
self.assertIn("table_name", call_args)
self.assertIn("table_type", call_args)
self.assertEqual(response, expected_response)

def test_database_name_is_bracket_quoted(self):
"""
Tests that the database name is bracket-quoted in catalog queries, so names with '-' or ']' work
"""
self.handler.database = "sqldb-prod]eastus"
self.handler.native_query = MagicMock(return_value=OkResponse())

for method in (
self.handler.get_tables,
self.handler.meta_get_tables,
self.handler.meta_get_columns,
self.handler.meta_get_column_statistics,
self.handler.meta_get_primary_keys,
self.handler.meta_get_foreign_keys,
):
with self.subTest(method=method.__name__):
self.handler.native_query.reset_mock()
method()
query = self.handler.native_query.call_args[0][0]
self.assertIn("[sqldb-prod]]eastus].", query)
self.assertEqual(query.count("[sqldb-prod]]eastus]."), query.count("sqldb-prod"))

def test_get_columns(self):
"""
Tests that get_columns calls native_query with the correct SQL
Expand Down
Loading