diff --git a/dcs_core/integrations/databases/mssql.py b/dcs_core/integrations/databases/mssql.py index 0d029a4c..c3dc58b4 100644 --- a/dcs_core/integrations/databases/mssql.py +++ b/dcs_core/integrations/databases/mssql.py @@ -126,20 +126,42 @@ def quote_column(self, column: str) -> str: return f"[{column}]" def query_get_table_names( - self, - schema: str | None = None, - ) -> List[str]: + self, schema: str | None = None, with_view: bool = False + ) -> dict: """ Get the list of tables in the database. :param schema: optional schema name - :return: list of table names + :param with_view: whether to include views + :return: dictionary with table names and optionally view names """ schema = schema or self.schema_name - query = f"SELECT o.name AS table_name FROM sys.objects o JOIN sys.schemas s ON o.schema_id = s.schema_id WHERE o.type = 'U' AND s.name = '{schema}' ORDER BY o.name" + + if with_view: + object_types = "IN ('U', 'V')" + else: + object_types = "= 'U'" + + query = f"SELECT o.name AS table_name, o.type FROM sys.objects o JOIN sys.schemas s ON o.schema_id = s.schema_id WHERE o.type {object_types} AND s.name = '{schema}' ORDER BY o.name" rows = self.fetchall(query) - res = [row[0] for row in rows] if rows else [] - return res + + if with_view: + result = {"table": [], "view": []} + if rows: + for row in rows: + object_name = row[0] + object_type = row[1].strip() if row[1] else row[1] + + if object_type == "U": + result["table"].append(object_name) + elif object_type == "V": + result["view"].append(object_name) + else: + result = {"table": []} + if rows: + result["table"] = [row[0] for row in rows] + + return result def query_get_table_columns( self, table: str, schema: str | None = None diff --git a/dcs_core/integrations/databases/mysql.py b/dcs_core/integrations/databases/mysql.py index 7ae5adb8..c4bc64fa 100644 --- a/dcs_core/integrations/databases/mysql.py +++ b/dcs_core/integrations/databases/mysql.py @@ -91,16 +91,40 @@ def quote_column(self, column: str) -> str: def query_get_table_names( self, schema: str | None = None, - ) -> List[str]: + with_view: bool = False, + ) -> dict: """ Get the list of tables in the database. :param schema: optional schema name - :return: list of table names + :param with_view: whether to include views + :return: dictionary with table names and optionally view names """ database = self.database - query = f"SELECT TABLES.TABLE_NAME FROM information_schema.tables WHERE TABLES.TABLE_SCHEMA = '{database}' and TABLES.TABLE_TYPE = 'BASE TABLE'" - result = self.fetchall(query) - return [row[0] for row in result] + if with_view: + table_type_condition = "TABLES.TABLE_TYPE IN ('BASE TABLE', 'VIEW')" + else: + table_type_condition = "TABLES.TABLE_TYPE = 'BASE TABLE'" + + query = f"SELECT TABLES.TABLE_NAME, TABLES.TABLE_TYPE FROM information_schema.tables WHERE TABLES.TABLE_SCHEMA = '{database}' and {table_type_condition}" + rows = self.fetchall(query) + + if with_view: + result = {"table": [], "view": []} + if rows: + for row in rows: + table_name = row[0] + table_type = row[1].strip() if row[1] else row[1] + + if table_type == "BASE TABLE": + result["table"].append(table_name) + elif table_type == "VIEW": + result["view"].append(table_name) + else: + result = {"table": []} + if rows: + result["table"] = [row[0] for row in rows] + + return result def query_get_table_columns( self, table: str, schema: str | None = None diff --git a/dcs_core/integrations/databases/oracle.py b/dcs_core/integrations/databases/oracle.py index 77cfefc9..28ccd21f 100644 --- a/dcs_core/integrations/databases/oracle.py +++ b/dcs_core/integrations/databases/oracle.py @@ -57,7 +57,9 @@ def connect(self) -> Any: "service_name": self.data_connection.get("service_name"), }, ) - self.schema_name = self.data_connection.get("username") + self.schema_name = self.data_connection.get( + "schema" + ) or self.data_connection.get("username") self.connection = engine.connect() return self.connection except Exception as e: @@ -86,16 +88,44 @@ def quote_column(self, column: str) -> str: def query_get_table_names( self, schema: str | None = None, - ) -> List[str]: + with_view: bool = False, + ) -> dict: """ Get the list of tables in the database. :param schema: optional schema name - :return: list of table names + :param with_view: whether to include views + :return: dictionary with table names and optionally view names """ schema = schema or self.schema_name - query = f"SELECT TABLE_NAME FROM ALL_ALL_TABLES WHERE OWNER = '{schema}'" - result = self.fetchall(query) - return [row[0] for row in result] + + if with_view: + query = ( + f"SELECT TABLE_NAME, 'TABLE' AS OBJECT_TYPE FROM ALL_ALL_TABLES WHERE OWNER = '{schema}' " + f"UNION " + f"SELECT VIEW_NAME AS TABLE_NAME, 'VIEW' AS OBJECT_TYPE FROM ALL_VIEWS WHERE OWNER = '{schema}'" + ) + else: + query = f"SELECT TABLE_NAME, 'TABLE' AS OBJECT_TYPE FROM ALL_ALL_TABLES WHERE OWNER = '{schema}'" + + rows = self.fetchall(query) + + if with_view: + result = {"table": [], "view": []} + if rows: + for row in rows: + object_name = row[0] + object_type = row[1].strip() if row[1] else row[1] + + if object_type == "TABLE": + result["table"].append(object_name) + elif object_type == "VIEW": + result["view"].append(object_name) + else: + result = {"table": []} + if rows: + result["table"] = [row[0] for row in rows] + + return result def query_get_table_columns( self, diff --git a/dcs_core/integrations/databases/postgres.py b/dcs_core/integrations/databases/postgres.py index 4d2fbd11..9e2618af 100644 --- a/dcs_core/integrations/databases/postgres.py +++ b/dcs_core/integrations/databases/postgres.py @@ -74,21 +74,46 @@ def quote_column(self, column: str) -> str: def query_get_table_names( self, schema: str | None = None, - ) -> List[str]: + with_view: bool = False, + ) -> dict: """ Get the list of tables in the database. :param schema: optional schema name - :return: list of table names + :param with_view: whether to include views + :return: dictionary with table names and optionally view names """ schema = schema or self.schema_name database = self.quote_database(self.database) + + if with_view: + table_type_condition = "table_type IN ('BASE TABLE', 'VIEW')" + else: + table_type_condition = "table_type = 'BASE TABLE'" + query = ( - f'SELECT table_name FROM {database}.information_schema.tables ' - f"WHERE table_schema = '{schema}' AND table_type = 'BASE TABLE'" + f"SELECT table_name, table_type FROM {database}.information_schema.tables " + f"WHERE table_schema = '{schema}' AND {table_type_condition}" ) - result = self.fetchall(query) - return [row[0] for row in result] + rows = self.fetchall(query) + + if with_view: + result = {"table": [], "view": []} + if rows: + for row in rows: + table_name = row[0] + table_type = row[1].strip() if row[1] else row[1] + + if table_type == "BASE TABLE": + result["table"].append(table_name) + elif table_type == "VIEW": + result["view"].append(table_name) + else: + result = {"table": []} + if rows: + result["table"] = [row[0] for row in rows] + + return result def query_get_table_columns( self, diff --git a/dcs_core/integrations/databases/sybase.py b/dcs_core/integrations/databases/sybase.py index 0bdfe1bd..29890def 100644 --- a/dcs_core/integrations/databases/sybase.py +++ b/dcs_core/integrations/databases/sybase.py @@ -331,27 +331,58 @@ def query_get_table_columns( def query_get_table_names( self, schema: str | None = None, - ) -> List[str]: + with_view: bool = False, + ) -> dict: """ Get the list of tables in the database. :param schema: optional schema name - :return: list of table names + :param with_view: whether to include views + :return: dictionary with table names and optionally view names """ schema = schema or self.schema_name database = self.database + if with_view: + type_condition = "IN ('U', 'V')" + else: + type_condition = "= 'U'" if self.sybase_driver_type.is_iq: - query = f"SELECT table_name FROM {database}.SYS.SYSTABLE WHERE creator = USER_ID('{schema}') AND table_type = 'BASE'" + table_type_condition = ( + "table_type IN ('BASE', 'VIEW')" if with_view else "table_type = 'BASE'" + ) + query = f"SELECT table_name, table_type FROM {database}.SYS.SYSTABLE WHERE creator = USER_ID('{schema}') AND {table_type_condition}" elif self.sybase_driver_type.is_ase: - query = f"SELECT name AS table_name FROM {database}..sysobjects WHERE type = 'U' AND uid = USER_ID('{schema}')" + query = f"SELECT name AS table_name, type FROM {database}..sysobjects WHERE type {type_condition} AND uid = USER_ID('{schema}')" elif self.sybase_driver_type.is_freetds: - query = f"SELECT name AS table_name FROM {database}.dbo.sysobjects WHERE type = 'U' AND uid = USER_ID('{schema}')" + query = f"SELECT name AS table_name, type FROM {database}.dbo.sysobjects WHERE type {type_condition} AND uid = USER_ID('{schema}')" else: raise ValueError("Unknown Sybase driver type") rows = self.fetchall(query) - res = [row[0] for row in rows] if rows else [] - return res + + if with_view: + result = {"table": [], "view": []} + if rows: + for row in rows: + table_name = row[0] + table_type = row[1].strip() if row[1] else row[1] + + if self.sybase_driver_type.is_iq: + if table_type == "BASE": + result["table"].append(table_name) + elif table_type == "VIEW": + result["view"].append(table_name) + else: # ASE or FreeTDS + if table_type == "U": + result["table"].append(table_name) + elif table_type == "V": + result["view"].append(table_name) + else: + result = {"table": []} + if rows: + result["table"] = [row[0] for row in rows] + + return result def convert_regex_to_sybase_pattern(self, regex_pattern: str) -> str: """