88
99from mindsdb_sql_parser import parse_sql
1010from mindsdb_sql_parser .ast .base import ASTNode
11+ from mindsdb_sql_parser .ast import Identifier
1112
1213from mindsdb .integrations .libs .base import MetaDatabaseHandler
14+ from mindsdb .integrations .utilities .query_traversal import query_traversal
1315from mindsdb .utilities import log
1416from mindsdb .utilities .render .sqlalchemy_render import SqlalchemyRender
1517from mindsdb .integrations .libs .response import (
@@ -169,6 +171,7 @@ def __init__(self, name, **kwargs):
169171 self .connection_args = kwargs .get ("connection_data" )
170172 self .dialect = "mssql"
171173 self .database = self .connection_args .get ("database" )
174+ self .schema = self .connection_args .get ("schema" )
172175 self .renderer = SqlalchemyRender ("mssql" )
173176
174177 # Determine if ODBC should be used
@@ -381,6 +384,26 @@ def native_query(self, query: str) -> Response:
381384
382385 return response
383386
387+ def _add_schema_to_tables (self , node , is_table = False , ** kwargs ):
388+ """
389+ Callback for query_traversal that adds schema prefix to table identifiers.
390+
391+ Args:
392+ node: The AST node being visited
393+ is_table: True if this node represents a table reference
394+ **kwargs: Other arguments from query_traversal (parent_query, callstack, etc.)
395+
396+ Returns:
397+ None to keep traversing, or a replacement node
398+ Note: This is mostly a workaround for Minds but it should still work for FQE
399+ """
400+ if is_table and isinstance (node , Identifier ):
401+ # Only add schema if the identifier doesn't already have one (single part)
402+ if len (node .parts ) == 1 :
403+ node .parts .insert (0 , self .schema )
404+ node .is_quoted .insert (0 , False )
405+ return None
406+
384407 def query (self , query : ASTNode ) -> Response :
385408 """
386409 Executes a SQL query represented by an ASTNode and retrieves the data.
@@ -391,6 +414,9 @@ def query(self, query: ASTNode) -> Response:
391414 Returns:
392415 Response: The response from the `native_query` method, containing the result of the SQL query execution.
393416 """
417+ # Add schema prefix to table identifiers if schema is configured
418+ if self .schema :
419+ query_traversal (query , self ._add_schema_to_tables )
394420
395421 query_str = self .renderer .get_string (query , with_failback = True )
396422 logger .debug (f"Executing SQL query: { query_str } )
@@ -410,8 +436,11 @@ def get_tables(self) -> Response:
410436 table_name,
411437 table_type
412438 FROM { self .database }
413- WHERE TABLE_TYPE in ('BASE TABLE', 'VIEW');
439+ WHERE TABLE_TYPE in ('BASE TABLE', 'VIEW')
414440 """
441+ if self .schema :
442+ query += f" AND table_schema = '{ self .schema }
443+
415444 return self .native_query (query )
416445
417446 def get_columns (self , table_name ) -> Response :
@@ -446,6 +475,10 @@ def get_columns(self, table_name) -> Response:
446475 WHERE
447476 table_name = '{ table_name }
448477 """
478+
479+ if self .schema :
480+ query += f" AND table_schema = '{ self .schema }
481+
449482 result = self .native_query (query )
450483 result .to_columns_table_response (map_type_fn = _map_type )
451484 return result
@@ -483,9 +516,13 @@ def meta_get_tables(self, table_names: list[str] | None = None) -> Response:
483516 AND p.index_id IN (0, 1)
484517 WHERE t.TABLE_TYPE IN ('BASE TABLE', 'VIEW')
485518 AND t.TABLE_SCHEMA NOT IN ('sys', 'INFORMATION_SCHEMA')
486- GROUP BY t.TABLE_NAME, t.TABLE_SCHEMA, t.TABLE_TYPE, ep.value
487519 """
488520
521+ if self .schema :
522+ query += f" AND t.TABLE_SCHEMA = '{ self .schema }
523+
524+ query += " GROUP BY t.TABLE_NAME, t.TABLE_SCHEMA, t.TABLE_TYPE, ep.value"
525+
489526 if table_names is not None and len (table_names ) > 0 :
490527 quoted_names = [f"'{ t } for t in table_names ]
491528 query += f" HAVING t.TABLE_NAME IN ({ ',' .join (quoted_names )}
@@ -525,6 +562,9 @@ def meta_get_columns(self, table_names: list[str] | None = None) -> Response:
525562 WHERE c.TABLE_SCHEMA NOT IN ('sys', 'INFORMATION_SCHEMA')
526563 """
527564
565+ if self .schema :
566+ query += f" AND c.TABLE_SCHEMA = '{ self .schema }
567+
528568 if table_names is not None and len (table_names ) > 0 :
529569 quoted_names = [f"'{ t } for t in table_names ]
530570 query += f" AND c.TABLE_NAME IN ({ ',' .join (quoted_names )}
@@ -553,6 +593,10 @@ def meta_get_column_statistics(self, table_names: list[str] | None = None) -> Re
553593 quoted_names = [f"'{ t } for t in table_names ]
554594 table_filter = f" AND t.name IN ({ ',' .join (quoted_names )}
555595
596+ schema_filter = ""
597+ if self .schema :
598+ schema_filter = f" AND s.name = '{ self .schema }
599+
556600 # Using OUTER APPLY to handle table-valued functions properly
557601 # This is equivalent to PostgreSQL's pg_stats view approach
558602 # Includes all statistics: auto-created, user-created, and index-based
@@ -589,6 +633,7 @@ def meta_get_column_statistics(self, table_names: list[str] | None = None) -> Re
589633 WHERE st.object_id IS NOT NULL
590634 ) h
591635 WHERE s.name NOT IN ('sys', 'INFORMATION_SCHEMA')
636+ { schema_filter }
592637 { table_filter }
593638 ORDER BY t.name, c.name
594639 """
@@ -620,6 +665,9 @@ def meta_get_primary_keys(self, table_names: list[str] | None = None) -> Respons
620665 WHERE tc.CONSTRAINT_TYPE = 'PRIMARY KEY'
621666 """
622667
668+ if self .schema :
669+ query += f" AND tc.TABLE_SCHEMA = '{ self .schema }
670+
623671 if table_names is not None and len (table_names ) > 0 :
624672 quoted_names = [f"'{ t } for t in table_names ]
625673 query += f" AND tc.TABLE_NAME IN ({ ',' .join (quoted_names )}
@@ -656,6 +704,9 @@ def meta_get_foreign_keys(self, table_names: list[str] | None = None) -> Respons
656704 WHERE s.name NOT IN ('sys', 'INFORMATION_SCHEMA')
657705 """
658706
707+ if self .schema :
708+ query += f" AND s.name = '{ self .schema }
709+
659710 if table_names is not None and len (table_names ) > 0 :
660711 quoted_names = [f"'{ t } for t in table_names ]
661712 query += f" AND OBJECT_NAME(fk.parent_object_id) IN ({ ',' .join (quoted_names )}
0 commit comments