MindsDB 1.3.1__py3-none-any.whl → 25.1.2.0__py3-none-any.whl

mindsdb/interfaces/storage/json.py

@@ -0,0 +1,97 @@
+from mindsdb.interfaces.storage import db
+from mindsdb.interfaces.storage.fs import RESOURCE_GROUP
+from mindsdb.utilities.context import context as ctx
+from mindsdb.utilities import log
+
+logger = log.getLogger(__name__)
+
+
+class JsonStorage:
+    def __init__(self, resource_group: str, resource_id: int):
+        self.resource_group = resource_group
+        self.resource_id = resource_id
+
+    def __setitem__(self, key, value):
+        if isinstance(value, dict) is False:
+            raise TypeError(f"got {type(value)} instead of dict")
+        existing_record = self.get_record(key)
+        if existing_record is None:
+            record = db.JsonStorage(
+                name=key,
+                resource_group=self.resource_group,
+                resource_id=self.resource_id,
+                company_id=ctx.company_id,
+                content=value
+            )
+            db.session.add(record)
+        else:
+            existing_record.content = value
+        db.session.commit()
+
+    def set(self, key, value):
+        self[key] = value
+
+    def __getitem__(self, key):
+        record = self.get_record(key)
+        if record is None:
+            return None
+        return record.content
+
+    def get(self, key):
+        return self[key]
+
+    def get_record(self, key):
+        record = db.session.query(db.JsonStorage).filter_by(
+            name=key,
+            resource_group=self.resource_group,
+            resource_id=self.resource_id,
+            company_id=ctx.company_id
+        ).first()
+        return record
+
+    def get_all_records(self):
+        records = db.session.query(db.JsonStorage).filter_by(
+            resource_group=self.resource_group,
+            resource_id=self.resource_id,
+            company_id=ctx.company_id
+        ).all()
+        return records
+
+    def __repr__(self):
+        records = self.get_all_records()
+        names = [x.name for x in records]
+        return f'json_storage({names})'
+
+    def __len__(self):
+        records = self.get_all_records()
+        return len(records)
+
+    def __delitem__(self, key):
+        record = self.get_record(key)
+        if record is not None:
+            try:
+                db.session.delete(record)
+                db.session.commit()
+            except Exception:
+                db.session.rollback()
+                logger.error('cant delete record from JSON storage')
+
+    def delete(self, key):
+        del self[key]
+
+    def clean(self):
+        json_records = self.get_all_records()
+        for record in json_records:
+            db.session.delete(record)
+        try:
+            db.session.commit()
+        except Exception:
+            db.session.rollback()
+            logger.error('cant delete records from JSON storage')
+
+
+def get_json_storage(resource_id: int, resource_group: str = RESOURCE_GROUP.PREDICTOR):
+    return JsonStorage(
+        resource_group=resource_group,
+        resource_id=resource_id,
+    )

mindsdb/interfaces/storage/model_fs.py

@@ -0,0 +1,282 @@
+import os
+import re
+import io
+import zipfile
+from typing import Union
+
+import mindsdb.interfaces.storage.db as db
+
+from .fs import RESOURCE_GROUP, FileStorageFactory, SERVICE_FILES_NAMES
+from .json import get_json_storage
+
+
+class ModelStorage:
+    """
+    This class deals with all model-related storage requirements, from setting status to storing artifacts.
+    """
+    def __init__(self, predictor_id):
+        storageFactory = FileStorageFactory(
+            resource_group=RESOURCE_GROUP.PREDICTOR,
+            sync=True
+        )
+        self.fileStorage = storageFactory(predictor_id)
+        self.predictor_id = predictor_id
+
+    # -- fields --
+
+    def _get_model_record(self, model_id: int, check_exists: bool = False) -> Union[db.Predictor, None]:
+        """Get model record by id
+
+        Args:
+            model_id (int): model id
+            check_exists (bool): true if need to check that model exists
+
+        Returns:
+            Union[db.Predictor, None]: model record
+
+        Raises:
+            KeyError: if `check_exists` is True and model does not exists
+        """
+        model_record = db.Predictor.query.get(self.predictor_id)
+        if check_exists is True and model_record is None:
+            raise KeyError('Model does not exists')
+        return model_record
+
+    def get_info(self):
+        rec = self._get_model_record(self.predictor_id)
+        return dict(status=rec.status,
+                    to_predict=rec.to_predict,
+                    data=rec.data,
+                    learn_args=rec.learn_args)
+
+    def status_set(self, status, status_info=None):
+        rec = self._get_model_record(self.predictor_id)
+        rec.status = status
+        if status_info is not None:
+            rec.data = status_info
+        db.session.commit()
+
+    def training_state_set(self, current_state_num=None, total_states=None, state_name=None):
+        rec = self._get_model_record(self.predictor_id)
+        if current_state_num is not None:
+            rec.training_phase_current = current_state_num
+        if total_states is not None:
+            rec.training_phase_total = total_states
+        if state_name is not None:
+            rec.training_phase_name = state_name
+        db.session.commit()
+
+    def training_state_get(self):
+        rec = self._get_model_record(self.predictor_id)
+        return [rec.training_phase_current, rec.training_phase_total, rec.training_phase_name]
+
+    def columns_get(self):
+        rec = self._get_model_record(self.predictor_id)
+        return rec.dtype_dict
+
+    def columns_set(self, columns):
+        # columns: {name: dtype}
+
+        rec = self._get_model_record(self.predictor_id)
+        rec.dtype_dict = columns
+        db.session.commit()
+
+    # files
+
+    def file_get(self, name):
+        return self.fileStorage.file_get(name)
+
+    def file_set(self, name, content):
+        self.fileStorage.file_set(name, content)
+
+    def folder_get(self, name):
+        # pull folder and return path
+        name = name.lower().replace(' ', '_')
+        name = re.sub(r'([^a-z^A-Z^_\d]+)', '_', name)
+
+        self.fileStorage.pull_path(name)
+        return str(self.fileStorage.get_path(name))
+
+    def folder_sync(self, name):
+        # sync abs path
+        name = name.lower().replace(' ', '_')
+        name = re.sub(r'([^a-z^A-Z^_\d]+)', '_', name)
+
+        self.fileStorage.push_path(name)
+
+    def file_list(self):
+        ...
+
+    def file_del(self, name):
+        ...
+
+    # jsons
+
+    def json_set(self, name, data):
+        json_storage = get_json_storage(
+            resource_id=self.predictor_id,
+            resource_group=RESOURCE_GROUP.PREDICTOR
+        )
+        return json_storage.set(name, data)
+
+    def json_get(self, name):
+        json_storage = get_json_storage(
+            resource_id=self.predictor_id,
+            resource_group=RESOURCE_GROUP.PREDICTOR
+        )
+        return json_storage.get(name)
+
+    def json_list(self):
+        ...
+
+    def json_del(self, name):
+        ...
+
+    def delete(self):
+        self.fileStorage.delete()
+        json_storage = get_json_storage(
+            resource_id=self.predictor_id,
+            resource_group=RESOURCE_GROUP.PREDICTOR
+        )
+        json_storage.clean()
+
+
+class HandlerStorage:
+    """
+    This class deals with all handler-related storage requirements, from storing metadata to synchronizing folders
+    across instances.
+    """
+    def __init__(self, integration_id: int, root_dir: str = None, is_temporal=False):
+        args = {}
+        if root_dir is not None:
+            args['root_dir'] = root_dir
+        storageFactory = FileStorageFactory(
+            resource_group=RESOURCE_GROUP.INTEGRATION,
+            sync=False,
+            **args
+        )
+        self.fileStorage = storageFactory(integration_id)
+        self.integration_id = integration_id
+        self.is_temporal = is_temporal
+        # do not sync with remote storage
+
+    def __convert_name(self, name):
+        name = name.lower().replace(' ', '_')
+        return re.sub(r'([^a-z^A-Z^_\d]+)', '_', name)
+
+    def is_empty(self):
+        """ check if storage directory is empty
+
+            Returns:
+                bool: true if dir is empty
+        """
+        for path in self.fileStorage.folder_path.iterdir():
+            if path.is_file() and path.name in SERVICE_FILES_NAMES:
+                continue
+            return False
+        return True
+
+    def get_connection_args(self):
+        rec = db.Integration.query.get(self.integration_id)
+        return rec.data
+
+    def update_connection_args(self, connection_args: dict) -> None:
+        """update integration connection args
+
+        Args:
+            connection_args (dict): new connection args
+        """
+        rec = db.Integration.query.get(self.integration_id)
+        if rec is None:
+            raise KeyError("Can't find integration")
+        rec.data = connection_args
+        db.session.commit()
+
+    # files
+
+    def file_get(self, name):
+        self.fileStorage.pull_path(name)
+        return self.fileStorage.file_get(name)
+
+    def file_set(self, name, content):
+        self.fileStorage.file_set(name, content)
+        if not self.is_temporal:
+            self.fileStorage.push_path(name)
+
+    def file_list(self):
+        ...
+
+    def file_del(self, name):
+        ...
+
+    # folder
+
+    def folder_get(self, name):
+        ''' Copies folder from remote to local file system and returns its path
+
+        :param name: name of the folder
+        '''
+        name = self.__convert_name(name)
+
+        self.fileStorage.pull_path(name)
+        return str(self.fileStorage.get_path(name))
+
+    def folder_sync(self, name):
+        # sync abs path
+        if self.is_temporal:
+            return
+        name = self.__convert_name(name)
+        self.fileStorage.push_path(name)
+
+    # jsons
+
+    def json_set(self, name, content):
+        json_storage = get_json_storage(
+            resource_id=self.integration_id,
+            resource_group=RESOURCE_GROUP.INTEGRATION
+        )
+        return json_storage.set(name, content)
+
+    def json_get(self, name):
+        json_storage = get_json_storage(
+            resource_id=self.integration_id,
+            resource_group=RESOURCE_GROUP.INTEGRATION
+        )
+        return json_storage.get(name)
+
+    def json_list(self):
+        ...
+
+    def json_del(self, name):
+        ...
+
+    def export_files(self) -> bytes:
+        if self.is_empty():
+            return None
+        folder_path = self.folder_get('')
+
+        zip_fd = io.BytesIO()
+
+        with zipfile.ZipFile(zip_fd, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for root, dirs, files in os.walk(folder_path):
+                for file_name in files:
+                    if file_name in SERVICE_FILES_NAMES:
+                        continue
+                    abs_path = os.path.join(root, file_name)
+                    zipf.write(abs_path, os.path.relpath(abs_path, folder_path))
+
+        zip_fd.seek(0)
+        return zip_fd.read()
+
+    def import_files(self, content: bytes):
+
+        folder_path = self.folder_get('')
+
+        zip_fd = io.BytesIO()
+        zip_fd.write(content)
+        zip_fd.seek(0)
+
+        with zipfile.ZipFile(zip_fd, 'r') as zip_ref:
+            zip_ref.extractall(folder_path)
+
+        self.folder_sync('')

mindsdb/interfaces/tabs/tabs_controller.py

@@ -0,0 +1,291 @@
+import json
+from typing import Dict, List
+from pathlib import Path
+
+from mindsdb.utilities import log
+from mindsdb.utilities.context import context as ctx
+from mindsdb.utilities.exception import EntityNotExistsError
+from mindsdb.interfaces.storage.fs import FileStorageFactory, RESOURCE_GROUP, FileStorage
+
+
+logger = log.getLogger(__name__)
+
+
+TABS_FILENAME = 'tabs'
+
+
+def get_storage():
+    # deprecated
+
+    storageFactory = FileStorageFactory(
+        resource_group=RESOURCE_GROUP.TAB,
+        sync=True
+    )
+
+    # resource_id is useless for 'tabs'
+    # use constant
+    return storageFactory(0)
+
+
+class TabsController:
+    """Tool for adding, editing, and deleting user's tabs
+
+    Attributes:
+        storage_factory (FileStorageFactory): callable object which returns tabs file storage
+    """
+
+    def __init__(self) -> None:
+        self.storage_factory = FileStorageFactory(
+            resource_group=RESOURCE_GROUP.TAB,
+            sync=True
+        )
+
+    def _get_file_storage(self) -> FileStorage:
+        """Get user's tabs file storage
+           NOTE: file storage depend is company_id sensitive, so need to recreate it each time
+
+        Returns:
+            FileStorage
+        """
+        return self.storage_factory(0)
+
+    def _get_next_tab_id(self) -> int:
+        """Get next free tab id
+
+        Returns:
+            int: id for next tab
+        """
+        tabs_files = self._get_tabs_files()
+        tabs_ids = list(tabs_files.keys())
+        if len(tabs_ids) == 0:
+            return 1
+        return max(tabs_ids) + 1
+
+    def _get_tabs_files(self) -> Dict[int, Path]:
+        """Get list of paths to each tab file
+
+        Returns:
+            Dict[int, Path]
+        """
+        tabs = {}
+        for child in self._get_file_storage().folder_path.iterdir():
+            if (child.is_file() and child.name.startswith('tab_')) is False:
+                continue
+            tab_id = child.name.replace('tab_', '')
+            if tab_id.isnumeric() is False:
+                continue
+            tabs[int(tab_id)] = child
+        return tabs
+
+    def _get_tabs_meta(self) -> List[Dict]:
+        """Get tabs info without content
+
+        Returns:
+            List[Dict]
+        """
+        all_tabs = self.get_all()
+        for tab in all_tabs:
+            del tab['content']
+        return all_tabs
+
+    def _migrate_legacy(self) -> None:
+        """Convert old single-file tabs storage to multiple files
+        """
+        file_storage = self._get_file_storage()
+        try:
+            file_data = file_storage.file_get(TABS_FILENAME)
+        except FileNotFoundError:
+            return
+        except Exception:
+            file_storage.delete()
+            return
+
+        try:
+            data = json.loads(file_data)
+        except Exception:
+            file_storage.delete()
+            return
+
+        if isinstance(data, dict) is False or isinstance(data.get('tabs'), str) is False:
+            file_storage.delete()
+            return
+
+        try:
+            tabs_list = json.loads(data['tabs'])
+        except Exception:
+            file_storage.delete()
+            return
+
+        if isinstance(tabs_list, list) is False:
+            file_storage.delete()
+            return
+
+        for tab in tabs_list:
+            tab_id = self._get_next_tab_id()
+
+            b_types = json.dumps({
+                'index': tab.get('index', 0),
+                'name': tab.get('name', 'undefined'),
+                'content': tab.get('value', '')
+            }).encode("utf-8")
+            file_storage.file_set(f'tab_{tab_id}', b_types)
+
+        file_storage.delete(TABS_FILENAME)
+
+    def get_all(self) -> List[Dict]:
+        """Get list of all tabs
+
+        Returns:
+            List[Dict]: all tabs data
+        """
+        self._get_file_storage().pull()
+        self._migrate_legacy()
+
+        tabs_files = self._get_tabs_files()
+        tabs_list = []
+        for tab_id, tab_path in tabs_files.items():
+            try:
+                data = json.loads(tab_path.read_text())
+            except Exception as e:
+                logger.error(f"Can't read data of tab {ctx.company_id}/{tab_id}: {e}")
+                continue
+            tabs_list.append({
+                'id': tab_id,
+                **data
+            })
+
+        tabs_list.sort(key=lambda x: x['index'])
+        return tabs_list
+
+    def get(self, tab_id: int) -> Dict:
+        """Get data of single tab
+
+        Args:
+            tab_id (int): id of the tab
+
+        Returns:
+            dict: tabs data
+        """
+        if isinstance(tab_id, int) is False:
+            raise ValueError('Tab id must be integer')
+
+        try:
+            raw_tab_data = self._get_file_storage().file_get(f'tab_{tab_id}')
+        except FileNotFoundError:
+            raise EntityNotExistsError(f'tab {tab_id}')
+
+        try:
+            data = json.loads(raw_tab_data)
+        except Exception as e:
+            logger.error(f"Can't read data of tab {ctx.company_id}/{tab_id}: {e}")
+            raise Exception(f"Can't read data of tab: {e}")
+
+        return {
+            'id': tab_id,
+            **data
+        }
+
+    def add(self, index: int = None, name: str = 'undefined', content: str = '') -> Dict:
+        """Add new tab
+
+        Args:
+            index (int, optional): index of new tab
+            name (str, optional): name of new tab
+            content (str, optional): content of new tab
+
+        Returns:
+            dict: new tab meta info: id, name and index
+        """
+        file_storage = self._get_file_storage()
+        tab_id = self._get_next_tab_id()
+
+        reorder_required = index is not None
+        if index is None:
+            all_tabs = self.get_all()
+            if len(all_tabs) == 0:
+                index = 0
+            else:
+                index = max([x.get('index', 0) for x in all_tabs]) + 1
+
+        data_bytes = json.dumps({
+            'index': index,
+            'name': name,
+            'content': content
+        }).encode("utf-8")
+        file_storage.file_set(f'tab_{tab_id}', data_bytes)
+
+        if reorder_required:
+            all_tabs = self.get_all()
+            all_tabs.sort(key=lambda x: (x['index'], 0 if x['id'] == tab_id else 1))
+            file_storage.sync = False
+            for tab_index, tab in enumerate(all_tabs):
+                tab['index'] = tab_index
+                data_bytes = json.dumps(tab).encode('utf-8')
+                file_storage.file_set(f'tab_{tab["id"]}', data_bytes)
+            file_storage.sync = True
+            file_storage.push()
+
+        return {'id': tab_id, 'index': index, 'name': name}
+
+    def modify(self, tab_id: int, index: int = None, name: str = None, content: str = None) -> Dict:
+        """Modify the tab
+
+        Args:
+            tab_id (int): if of the tab to modify
+            index (int, optional): tab's new index
+            name (str, optional): tab's new name
+            content (str, optional): tab's new content
+
+        Returns:
+            dict: new tab meta info: id, name and index
+        """
+        file_storage = self._get_file_storage()
+        current_data = self.get(tab_id)
+
+        # region modify index
+        if index is not None and current_data['index'] != index:
+            current_data['index'] = index
+            all_tabs = [x for x in self.get_all() if x['id'] != tab_id]
+            all_tabs.sort(key=lambda x: x['index'])
+            file_storage.sync = False
+            for tab_index, tab in enumerate(all_tabs):
+                if tab_index < index:
+                    tab['index'] = tab_index
+                else:
+                    tab['index'] = tab_index + 1
+                data_bytes = json.dumps(tab).encode('utf-8')
+                file_storage.file_set(f'tab_{tab["id"]}', data_bytes)
+            file_storage.sync = True
+            file_storage.push()
+        # endregion
+
+        # region modify name
+        if name is not None and current_data['name'] != name:
+            current_data['name'] = name
+        # endregion
+
+        # region modify content
+        if content is not None and current_data['content'] != content:
+            current_data['content'] = content
+        # endregion
+
+        data_bytes = json.dumps(current_data).encode('utf-8')
+        file_storage.file_set(f'tab_{tab_id}', data_bytes)
+
+        return {
+            'id': current_data['id'],
+            'index': current_data['index'],
+            'name': current_data['name']
+        }
+
+    def delete(self, tab_id: int):
+        file_storage = self._get_file_storage()
+        try:
+            file_storage.file_get(f'tab_{tab_id}')
+        except FileNotFoundError:
+            raise EntityNotExistsError(f'tab {tab_id}')
+
+        file_storage.delete(f'tab_{tab_id}')
+
+
+tabs_controller = TabsController()

mindsdb/interfaces/tasks/task.py

@@ -0,0 +1,16 @@
+from mindsdb.interfaces.storage import db
+
+
+class BaseTask:
+
+    def __init__(self, task_id, object_id):
+        self.task_id = task_id
+        self.object_id = object_id
+
+    def run(self, stop_event):
+        raise NotImplementedError
+
+    def set_error(self, message):
+        task_record = db.Tasks.query.get(self.task_id)
+        task_record.last_error = str(message)
+        db.session.commit()