add: view and table return to query_get_table_names

Author: AnshumanX304

dcs_core/integrations/databases/mssql.py

@@ -127,11 +127,12 @@ def quote_column(self, column: str) -> str:
 
     def query_get_table_names(
         self, schema: str | None = None, with_view: bool = False
-    ) -> List[str]:
+    ) -> 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
 
@@ -140,11 +141,27 @@ def query_get_table_names(
         else:
             object_types = "= 'U'"
 
-        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 {object_types} AND s.name = '{schema}' ORDER BY o.name"
+        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

dcs_core/integrations/databases/mysql.py

@@ -92,21 +92,39 @@ def query_get_table_names(
         self,
         schema: str | None = None,
         with_view: bool = False,
-    ) -> List[str]:
+    ) -> 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
         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 FROM information_schema.tables WHERE TABLES.TABLE_SCHEMA = '{database}' and {table_type_condition}"
-        result = self.fetchall(query)
-        return [row[0] for row in result]
+        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

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:
@@ -87,23 +89,43 @@ def query_get_table_names(
         self,
         schema: str | None = None,
         with_view: bool = False,
-    ) -> List[str]:
+    ) -> 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
+
         if with_view:
             query = (
-                f"SELECT TABLE_NAME FROM ALL_ALL_TABLES WHERE OWNER = '{schema}' "
+                f"SELECT TABLE_NAME, 'TABLE' AS OBJECT_TYPE FROM ALL_ALL_TABLES WHERE OWNER = '{schema}' "
                 f"UNION "
-                f"SELECT VIEW_NAME AS TABLE_NAME FROM ALL_VIEWS WHERE OWNER = '{schema}'"
+                f"SELECT VIEW_NAME AS TABLE_NAME, 'VIEW' AS OBJECT_TYPE FROM ALL_VIEWS WHERE OWNER = '{schema}'"
             )
         else:
-            query = f"SELECT TABLE_NAME FROM ALL_ALL_TABLES WHERE OWNER = '{schema}'"
-        result = self.fetchall(query)
-        return [row[0] for row in result]
+            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,

dcs_core/integrations/databases/postgres.py

@@ -75,11 +75,12 @@ def query_get_table_names(
         self,
         schema: str | None = None,
         with_view: bool = False,
-    ) -> List[str]:
+    ) -> 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
@@ -91,11 +92,28 @@ def query_get_table_names(
             table_type_condition = "table_type = 'BASE TABLE'"
 
         query = (
-            f"SELECT table_name FROM {database}.information_schema.tables "
+            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,

dcs_core/integrations/databases/sybase.py

@@ -332,11 +332,12 @@ def query_get_table_names(
         self,
         schema: str | None = None,
         with_view: bool = False,
-    ) -> List[str]:
+    ) -> 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
@@ -349,17 +350,39 @@ def query_get_table_names(
             table_type_condition = (
                 "table_type IN ('BASE', 'VIEW')" if with_view else "table_type = 'BASE'"
             )
-            query = f"SELECT table_name FROM {database}.SYS.SYSTABLE WHERE creator = USER_ID('{schema}') AND {table_type_condition}"
+            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 {type_condition} 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 {type_condition} 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:
         """