create job with sdk commands

Author: ea-rus

mindsdb_sdk/jobs.py

@@ -1,23 +1,32 @@
 import datetime as dt
 from typing import Union, List
 
+
 import pandas as pd
 
 from mindsdb_sql.parser.dialects.mindsdb import CreateJob, DropJob
 from mindsdb_sql.parser.ast import Identifier, Star, Select
 
+from mindsdb_sdk.query import Query
 from mindsdb_sdk.utils.sql import dict_to_binary_op
 from mindsdb_sdk.utils.objects_collection import CollectionBase
+from mindsdb_sdk.utils.context import set_saving
 
 
 class Job:
-    def __init__(self, project, data):
+    def __init__(self, project, name, data=None, create_callback=None):
         self.project = project
+        self.name = name
         self.data = data
-        self._update(data)
+
+        self.query_str = None
+        if data is not None:
+            self._update(data)
+        self._queries = []
+        self._create_callback = create_callback
 
     def _update(self, data):
-        self.name = data['name']
+        # self.name = data['name']
         self.query_str = data['query']
         self.start_at = data['start_at']
         self.end_at = data['end_at']
@@ -27,13 +36,52 @@ def _update(self, data):
     def __repr__(self):
         return f"{self.__class__.__name__}({self.name}, query='{self.query_str}')"
 
+    def __enter__(self):
+        if self._create_callback is None:
+            raise RuntimeError('It can not be used to create context')
+        set_saving(f'job-{self.name}')
+
+        return self
+
+    def __exit__(self, type, value, traceback):
+        set_saving(None)
+        if type is None:
+            if len(self._queries) == 0:
+                raise RuntimeError('No queries were added to job')
+
+            query_str = '; '.join(self._queries)
+
+            self._create_callback(query_str)
+
+            self.refresh()
+
     def refresh(self):
         """
         Retrieve job data from mindsdb server
         """
         job = self.project.get_job(self.name)
         self._update(job.data)
 
+    def add_query(self, query: Union[Query, str]):
+        """
+        Add a query to job. Method is used in context of the job
+
+        >>> with con.jobs.create('j1') as job:
+        >>>    job.add_query(table1.insert(table2))
+
+        :param query: string or Query object. Query.database should be emtpy or the same as job's project
+        """
+        if isinstance(query, Query):
+
+            if query.database is not None and query.database != self.project.name:
+                # we can't execute this query in jobs project
+                raise ValueError(f"Wrong query database: {query.database}. You could try to use sql string instead")
+
+            query = query.sql
+        elif not isinstance(query, str):
+            raise RuntimeError(f'Unable to use add this object as a query: {query}. Try to use sql string instead')
+        self._queries.append(query)
+
     def get_history(self) -> pd.DataFrame:
         """
         Get history of job execution
@@ -69,7 +117,7 @@ def _list(self, name: str = None) -> List[Job]:
         df = df.rename(columns=cols_map)
 
         return [
-            Job(self.project, item)
+            Job(self.project, item.pop('name'), item)
             for item in df.to_dict('records')
         ]
 
@@ -101,7 +149,7 @@ def get(self, name: str) -> Job:
     def create(
             self,
             name: str,
-            query_str: str,
+            query_str: str = None,
             start_at: dt.datetime = None,
             end_at: dt.datetime = None,
             repeat_str: str = None,
@@ -113,7 +161,25 @@ def create(
         If it is not possible (job executed and not accessible anymore):
            return None
 
-        More info: https://docs.mindsdb.com/sql/create/jobs
+        Usage options:
+
+        Option 1: to use string query
+        All job tasks could be passed as string with sql queries. Job is created emmideiately
+
+        >>> job = con.jobs.create('j1', query_str='retrain m1; show models', repeat_min=1):
+
+        Option 2: to use 'with' block.
+        It allows to pass sdk commands to job tasks.
+        Not all sdk commands could be accepted here,
+         only those which are converted in to sql in sdk and sent to /query endpoint
+        Adding query sql string is accepted as well
+        Job will be created after exit from 'with' block
+
+        >>> with con.jobs.create('j1', repeat_min=1) as job:
+        >>>    job.add_query(table1.insert(table2))
+        >>>    job.add_query('retrain m1')  # using string
+
+        More info about jobs: https://docs.mindsdb.com/sql/create/jobs
 
         :param name: name of the job
         :param query_str: str, job's query (or list of queries with ';' delimiter) which job have to execute
@@ -137,20 +203,30 @@ def create(
         if repeat_min is not None:
             repeat_str = f'{repeat_min} minutes'
 
-        ast_query = CreateJob(
-            name=Identifier(name),
-            query_str=query_str,
-            start_str=start_str,
-            end_str=end_str,
-            repeat_str=repeat_str
-        )
+        def _create_callback(query):
+            ast_query = CreateJob(
+                name=Identifier(name),
+                query_str=query,
+                start_str=start_str,
+                end_str=end_str,
+                repeat_str=repeat_str
+            )
+
+            self.api.sql_query(ast_query.to_string(), database=self.project.name)
+
+        if query_str is None:
+            # allow to create context with job
+            job = Job(self.project, name, create_callback=_create_callback)
+            return job
+        else:
+            # create it
+            _create_callback(query_str)
 
-        self.api.sql_query(ast_query.to_string(), database=self.project.name)
+            # job can be executed and remove it is not repeatable
+            jobs = self._list(name)
+            if len(jobs) == 1:
+                return jobs[0]
 
-        # job can be executed and remove it is not repeatable
-        jobs = self._list(name)
-        if len(jobs) == 1:
-            return jobs[0]
 
     def drop(self, name: str):
         """

mindsdb_sdk/models.py

@@ -15,6 +15,7 @@
 
 from mindsdb_sdk.utils.objects_collection import CollectionBase
 from mindsdb_sdk.utils.sql import dict_to_binary_op, query_to_native_query
+from mindsdb_sdk.utils.context import is_saving
 
 from .query import Query
 
@@ -106,15 +107,15 @@ def __repr__(self):
         version = ''
         if self.version is not None:
             version = f', version={self.version}'
-        return f'{self.__class__.__name__}({self.name}{version}, status={self.data["status"]})'
+        return f'{self.__class__.__name__}({self.name}{version}, status={self.data.get("status")})'
 
     def _get_identifier(self):
         parts = [self.project.name, self.name]
         if self.version is not None:
             parts.append(str(self.version))
         return Identifier(parts=parts)
 
-    def predict(self, data: Union[pd.DataFrame, Query, dict], params: dict = None) -> pd.DataFrame:
+    def predict(self, data: Union[pd.DataFrame, Query, dict], params: dict = None) -> Union[pd.DataFrame, Query]:
         """
         Make prediction using model
 
@@ -203,7 +204,11 @@ def predict(self, data: Union[pd.DataFrame, Query, dict], params: dict = None) -
             if params is not None:
                 upper_query.using = params
             # execute in query's database
-            return self.project.api.sql_query(upper_query.to_string(), database=None)
+            sql = upper_query.to_string()
+            if is_saving():
+                return Query(self, sql)
+
+            return self.project.api.sql_query(sql, database=None)
 
         elif isinstance(data, dict):
             data = pd.DataFrame([data])
@@ -310,15 +315,19 @@ def _retrain(self,
             integration_name=database,
             using=options or None,
         )
+        sql = ast_query.to_string()
+
+        if is_saving():
+            return Query(self, sql)
 
-        data = self.project.query(ast_query.to_string()).fetch()
+        data = self.project.api.sql_query(sql)
         data = {k.lower(): v for k, v in data.items()}
 
         # return new instance
         base_class = self.__class__
         return base_class(self.project, data)
 
-    def describe(self, type: str = None) -> pd.DataFrame:
+    def describe(self, type: str = None) -> Union[pd.DataFrame, Query]:
         """
         Return description of the model
 
@@ -332,7 +341,12 @@ def describe(self, type: str = None) -> pd.DataFrame:
         if type is not None:
             identifier.parts.append(type)
         ast_query = Describe(identifier)
-        return self.project.query(ast_query.to_string()).fetch()
+
+        sql = ast_query.to_string()
+        if is_saving():
+            return Query(self, sql)
+
+        return self.project.api.sql_query(sql)
 
     def list_versions(self) -> List[ModelVersion]:
         """
@@ -374,7 +388,7 @@ def set_active(self, version: int):
         :param version: version to set active
         """
         ast_query = Update(
-            table=Identifier('models_versions'),
+            table=Identifier(parts=[self.project.name, 'models_versions']),
             update_columns={
                 'active': Constant(1)
             },
@@ -383,7 +397,11 @@ def set_active(self, version: int):
                 'version': version
             })
         )
-        self.project.query(ast_query.to_string()).fetch()
+        sql = ast_query.to_string()
+        if is_saving():
+            return Query(self, sql)
+
+        self.project.api.sql_query(sql)
         self.refresh()
 
 
@@ -430,7 +448,7 @@ def create(
         database: str = None,
         options: dict = None,
         timeseries_options: dict = None, **kwargs
-    ) -> Model:
+    ) -> Union[Model, Query]:
         """
         Create new model in project and return it
 
@@ -486,7 +504,7 @@ def create(
             targets = None
 
         ast_query = CreatePredictor(
-            name=Identifier(name),
+            name=Identifier(parts=[self.project.name, name]),
             query_str=query,
             integration_name=database,
             targets=targets,
@@ -522,7 +540,13 @@ def create(
 
             options['engine'] = engine
         ast_query.using = options
-        df = self.project.query(ast_query.to_string()).fetch()
+
+        sql = ast_query.to_string()
+
+        if is_saving():
+            return Query(self, sql)
+
+        df = self.project.api.sql_query(sql)
         if len(df) > 0:
             data = dict(df.iloc[0])
             # to lowercase
@@ -559,8 +583,12 @@ def drop(self, name: str):
 
         :param name: name of the model
         """
-        ast_query = DropPredictor(name=Identifier(name))
-        self.project.query(ast_query.to_string()).fetch()
+        ast_query = DropPredictor(name=Identifier(parts=[self.project.name, name]))
+        sql = ast_query.to_string()
+        if is_saving():
+            return Query(self, sql)
+
+        self.project.api.sql_query(sql)
 
 
     def list(self, with_versions: bool = False,

mindsdb_sdk/tables.py

@@ -9,6 +9,7 @@
 
 from mindsdb_sdk.utils.sql import dict_to_binary_op, add_condition, query_to_native_query
 from mindsdb_sdk.utils.objects_collection import CollectionBase
+from mindsdb_sdk.utils.context import is_saving
 
 from .query import Query
 
@@ -159,7 +160,11 @@ def delete(self, **kwargs):
             where=dict_to_binary_op(kwargs)
         )
         sql = ast_query.to_string()
-        self.api.sql_query(sql, 'mindsdb')
+
+        if is_saving():
+            return Query(self, sql)
+
+        self.api.sql_query(sql)
 
     def update(self, values: Union[dict, Query], on: list = None, filters: dict = None):
         '''
@@ -218,6 +223,11 @@ def update(self, values: Union[dict, Query], on: list = None, filters: dict = No
         else:
             raise NotImplementedError
 
+        if is_saving():
+            return Query(self, sql)
+
+        self.api.sql_query(sql)
+
 
 class Tables(CollectionBase):
     """

mindsdb_sdk/utils/context.py

@@ -0,0 +1,25 @@
+from contextvars import ContextVar
+
+context_storage = ContextVar('create_context')
+
+
+def set_context(name, value):
+    data = context_storage.get({})
+    data[name] = value
+
+    context_storage.set(data)
+
+
+def get_context(name):
+
+    data = context_storage.get({})
+    return data.get(name)
+
+
+def set_saving(name):
+    set_context('saving', name)
+
+
+def is_saving():
+    return get_context('saving') is not None
+