diff --git a/mindsdb_sdk/__init__.py b/mindsdb_sdk/__init__.py index df54ae1..82a757c 100644 --- a/mindsdb_sdk/__init__.py +++ b/mindsdb_sdk/__init__.py @@ -1 +1,2 @@ from mindsdb_sdk.connect import connect +from mindsdb_sdk.tree import TreeNode diff --git a/mindsdb_sdk/connectors/rest_api.py b/mindsdb_sdk/connectors/rest_api.py index 4f3e36a..4b37efd 100644 --- a/mindsdb_sdk/connectors/rest_api.py +++ b/mindsdb_sdk/connectors/rest_api.py @@ -132,8 +132,12 @@ 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()) diff --git a/mindsdb_sdk/databases.py b/mindsdb_sdk/databases.py index 368ddf0..011be4f 100644 --- a/mindsdb_sdk/databases.py +++ b/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 @@ -54,6 +55,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}") + """ + df = self.api.objects_tree(self.name, with_schemas=with_schemas) + return [TreeNode.from_dict(row.to_dict()) for _, row in df.iterrows()] class Databases(CollectionBase): diff --git a/mindsdb_sdk/server.py b/mindsdb_sdk/server.py index e953c26..80a4bf4 100644 --- a/mindsdb_sdk/server.py +++ b/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 from .config import Config @@ -61,6 +64,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}") + """ + df = self.api.objects_tree('') + return [TreeNode.from_dict(row.to_dict()) for _, row in df.iterrows()] def __repr__(self): return f'{self.__class__.__name__}({self.api.url})' diff --git a/mindsdb_sdk/tree.py b/mindsdb_sdk/tree.py new file mode 100644 index 0000000..5f41a5d --- /dev/null +++ b/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