Merge pull request #101 from mindsdb/jobs

Job supports SDK comands

Author: ea-rus

mindsdb_sdk/connectors/rest_api.py

@@ -5,7 +5,7 @@
 import requests
 import pandas as pd
 
-from .. import __about__
+from mindsdb_sdk import __about__
 
 
 def _try_relogin(fnc):
@@ -73,7 +73,11 @@ def login(self):
         _raise_for_status(r)
 
     @_try_relogin
-    def sql_query(self, sql, database='mindsdb', lowercase_columns=False):
+    def sql_query(self, sql, database=None, lowercase_columns=False):
+
+        if database is None:
+            # it means the database is included in query
+            database = 'mindsdb'
         url = self.url + '/api/sql/query'
         r = self.session.post(url, json={
             'query': sql,

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,55 @@ 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 ValueError("The job is already created and can't be used to create context."
+                               " To be able to use context: create job without 'query_str' parameter: "
+                               "\n>>> with con.jobs.create('j1') as job:"
+                               "\n>>>    job.add_query(...)")
+        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 ValueError('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 ValueError(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 +120,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 +152,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 +164,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 immediately
+
+        >>> 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 +206,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/knowledge_bases.py

@@ -7,8 +7,9 @@
 from mindsdb_sql.parser.dialects.mindsdb import CreateKnowledgeBase, DropKnowledgeBase
 from mindsdb_sql.parser.ast import Identifier, Star, Select, BinaryOperation, Constant, Insert
 
-from mindsdb_sdk.utils.sql import dict_to_binary_op
+from mindsdb_sdk.utils.sql import dict_to_binary_op, query_to_native_query
 from mindsdb_sdk.utils.objects_collection import CollectionBase
+from mindsdb_sdk.utils.context import is_saving
 
 from .models import Model
 from .tables import Table
@@ -35,6 +36,7 @@ def __init__(self, project, data: dict):
 
         self.project = project
         self.name = data['name']
+        self.table_name = Identifier(parts=[self.project.name, self.name])
 
         self.storage = None
         if data['storage'] is not None:
@@ -69,10 +71,10 @@ def __init__(self, project, data: dict):
         self._query = None
         self._limit = None
 
-        database = project.name
         self._update_query()
 
-        super().__init__(project.api, self.sql, database)
+        # empty database
+        super().__init__(project.api, self.sql, None)
 
     def __repr__(self):
         return f'{self.__class__.__name__}({self.project.name}.{self.name})'
@@ -104,9 +106,7 @@ def _update_query(self):
 
         ast_query = Select(
             targets=[Star()],
-            from_table=Identifier(parts=[
-                self.project.name, self.name
-            ])
+            from_table=self.table_name
         )
         if self._query is not None:
             ast_query.where = BinaryOperation(op='=', args=[
@@ -143,20 +143,28 @@ def insert(self, data: Union[pd.DataFrame, Query, dict]):
             data_split = data.to_dict('split')
 
             ast_query = Insert(
-                table=Identifier(self.name),
+                table=Identifier(self.table_name),
                 columns=data_split['columns'],
                 values=data_split['data']
             )
-
             sql = ast_query.to_string()
-            self.api.sql_query(sql, self.database)
+
         else:
             # insert from select
-            table = Identifier(parts=[self.database, self.name])
-            self.api.sql_query(
-                f'INSERT INTO {table.to_string()} ({data.sql})',
-                database=data.database
-            )
+            if data.database is not None:
+                ast_query = Insert(
+                    table=Identifier(self.table_name),
+                    from_select=query_to_native_query(data)
+                )
+                sql = ast_query.to_string()
+            else:
+                sql = f'INSERT INTO {self.table_name.to_string()} ({data.sql})'
+
+        if is_saving():
+            # don't execute it right now, return query object
+            return Query(self, sql, self.database)
+
+        self.api.sql_query(sql, self.database)
 
 
 class KnowledgeBases(CollectionBase):
@@ -241,7 +249,7 @@ def create(
         content_columns: list = None,
         id_column: str = None,
         params: dict = None,
-    ) -> KnowledgeBase:
+    ) -> Union[KnowledgeBase, Query]:
         """
 
         Create knowledge base:
@@ -291,13 +299,17 @@ def create(
             storage_name = None
 
         ast_query = CreateKnowledgeBase(
-            Identifier(name),
+            Identifier(parts=[self.project.name, name]),
             model=model_name,
             storage=storage_name,
             params=params_out
         )
+        sql = ast_query.to_string()
 
-        self.api.sql_query(ast_query.to_string(), database=self.project.name)
+        if is_saving():
+            return Query(self, sql)
+
+        self.api.sql_query(sql)
 
         return self.get(name)
 
@@ -308,6 +320,10 @@ def drop(self, name: str):
         :return:
         """
 
-        ast_query = DropKnowledgeBase(Identifier(name))
+        ast_query = DropKnowledgeBase(Identifier(parts=[self.project.name, name]))
+        sql = ast_query.to_string()
+
+        if is_saving():
+            return Query(self, sql)
 
-        self.api.sql_query(ast_query.to_string())
+        self.api.sql_query(sql)