Add tree API

Author: ZoranPandovski

mindsdb_sdk/__init__.py

@@ -1 +1,2 @@
 from mindsdb_sdk.connect import connect
+from mindsdb_sdk.tree import TreeNode

mindsdb_sdk/connectors/rest_api.py

@@ -128,11 +128,15 @@ def model_predict(self, project, model, data, params=None, version=None):
         return pd.DataFrame(r.json())
 
     @_try_relogin
-    def objects_tree(self, item=''):
-        r = self.session.get(self.url + f'/api/tree/{item}')
+    def objects_tree(self, item='', with_schemas=False):
+        params = {}
+        if with_schemas:
+            params['all_schemas'] = 'true'
+        
+        r = self.session.get(self.url + f'/api/tree/{item}', params=params)
         _raise_for_status(r)
 
-        return pd.DataFrame(r.json())
+        return r.json()  # Return raw JSON instead of DataFrame for more flexible processing
 
     @staticmethod
     def read_file_as_bytes(file_path: str):

mindsdb_sdk/databases.py

@@ -4,6 +4,7 @@
 from mindsdb_sql.parser.ast import DropDatabase, Identifier
 
 from mindsdb_sdk.utils.objects_collection import CollectionBase
+from .tree import TreeNode
 
 from .query import Query
 from .tables import Tables
@@ -52,6 +53,35 @@ def query(self, sql: str) -> Query:
         :return: Query object
         """
         return Query(self.api, sql, database=self.name)
+    
+    def tree(self, with_schemas: bool = False) -> List[TreeNode]:
+        """
+        Get the tree structure of tables and schemas within this database.
+        
+        This returns a list of table/schema nodes with their metadata including:
+        - name: table/schema name
+        - class: node type ('table', 'schema', 'job')
+        - type: table type ('table', 'view', 'job', 'system view')
+        - engine: table engine (if applicable)
+        - deletable: whether the item can be deleted
+        - schema: schema name (for tables)
+        - children: nested tables (for schemas)
+        
+        :param with_schemas: Whether to include schema information for data databases
+        :return: List of TreeNode objects representing tables/schemas
+        
+        Example:
+        
+        >>> db = server.databases.get('my_postgres')
+        >>> tree = db.tree(with_schemas=True)
+        >>> for item in tree:
+        ...     if item.class_ == 'schema':
+        ...         print(f"Schema: {item.name} with {len(item.children)} tables")
+        ...     else:
+        ...         print(f"Table: {item.name}, Type: {item.type}")
+        """
+        tree_data = self.api.objects_tree(self.name, with_schemas=with_schemas)
+        return [TreeNode.from_dict(item) for item in tree_data]
 
 
 class Databases(CollectionBase):

mindsdb_sdk/server.py

@@ -1,9 +1,12 @@
+from typing import List
+
 from .agents import Agents
 from .databases import Databases
 from .projects import Project, Projects
 from .ml_engines import MLEngines
 from .handlers import Handlers
 from .skills import Skills
+from .tree import TreeNode
 
 
 class Server(Project):
@@ -58,6 +61,29 @@ def status(self) -> dict:
         :return: server status info
         """
         return self.api.status()
+    
+    def tree(self) -> List[TreeNode]:
+        """
+        Get the tree structure of databases on the server.
+        
+        This returns a list of database nodes with their metadata including:
+        - name: database name
+        - class: node type ('db')
+        - type: database type ('data', 'project', 'system')
+        - engine: database engine
+        - deletable: whether the database can be deleted
+        - visible: whether the database is visible
+        
+        :return: List of TreeNode objects representing databases
+        
+        Example:
+        
+        >>> tree = server.tree()
+        >>> for db in tree:
+        ...     print(f"Database: {db.name}, Type: {db.type}, Engine: {db.engine}")
+        """
+        tree_data = self.api.objects_tree('')
+        return [TreeNode.from_dict(item) for item in tree_data]
 
     def __repr__(self):
         return f'{self.__class__.__name__}({self.api.url})'

mindsdb_sdk/tree.py

@@ -0,0 +1,57 @@
+from typing import List, Optional, Dict, Any
+from dataclasses import dataclass
+
+
+@dataclass
+class TreeNode:
+    """Represents any node in the MindsDB tree structure."""
+    name: str
+    class_: str  # 'class' is a reserved keyword, so using 'class_'
+    type: Optional[str] = None
+    engine: Optional[str] = None
+    deletable: bool = False
+    visible: bool = True
+    schema: Optional[str] = None  # For table nodes that have schema information
+    children: Optional[List['TreeNode']] = None
+    
+    def __post_init__(self):
+        if self.children is None:
+            self.children = []
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'TreeNode':
+        """Create TreeNode from dictionary data."""
+        children = []
+        if 'children' in data and data['children']:
+            children = [cls.from_dict(child) for child in data['children']]
+        
+        return cls(
+            name=data['name'],
+            class_=data.get('class', ''),
+            type=data.get('type'),
+            engine=data.get('engine'),
+            deletable=data.get('deletable', False),
+            visible=data.get('visible', True),
+            schema=data.get('schema'),  # Include schema if present
+            children=children
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert TreeNode to dictionary."""
+        result = {
+            'name': self.name,
+            'class': self.class_,
+            'deletable': self.deletable,
+            'visible': self.visible
+        }
+        
+        if self.type is not None:
+            result['type'] = self.type
+        if self.engine is not None:
+            result['engine'] = self.engine
+        if self.schema is not None:
+            result['schema'] = self.schema
+        if self.children:
+            result['children'] = [child.to_dict() for child in self.children]
+            
+        return result