Merge pull request #44 from mindsdb/upsert

Upsert minds/datasources

Author: ea-rus

minds/datasources/datasources.py

@@ -16,11 +16,12 @@ class DatabaseConfig(BaseModel):
 class Datasource(DatabaseConfig):
     ...
 
+
 class Datasources:
     def __init__(self, client):
         self.api = client.api
 
-    def create(self, ds_config: DatabaseConfig, replace=False):
+    def create(self, ds_config: DatabaseConfig, update=False):
         """
         Create new datasource and return it
 
@@ -30,19 +31,16 @@ def create(self, ds_config: DatabaseConfig, replace=False):
            - description: str, description of the database. Used by mind to know what data can be got from it.
            - connection_data: dict, optional, credentials to connect to database
            - tables: list of str, optional, list of allowed tables
+        :param update: if true - to update datasourse if exists, default is false
         :return: datasource object
         """
 
         name = ds_config.name
 
-        if replace:
-            try:
-                self.get(name)
-                self.drop(name, force=True)
-            except exc.ObjectNotFound:
-                ...
-
-        self.api.post('/datasources', data=ds_config.model_dump())
+        if update:
+            self.api.put('/datasources', data=ds_config.model_dump())
+        else:
+            self.api.post('/datasources', data=ds_config.model_dump())
         return self.get(name)
 
     def list(self) -> List[Datasource]:

minds/minds.py

@@ -243,6 +243,7 @@ def create(
         datasources=None,
         parameters=None,
         replace=False,
+        update=False,
     ) -> Mind:
         """
         Create a new mind and return it
@@ -259,6 +260,7 @@ def create(
         :param datasources: list of datasources used by mind, optional
         :param parameters, dict: other parameters of the mind, optional
         :param replace: if true - to remove existing mind, default is false
+        :param update: if true - to update mind if exists, default is false
         :return: created mind
         """
 
@@ -284,7 +286,12 @@ def create(
         if 'prompt_template' not in parameters:
             parameters['prompt_template'] = DEFAULT_PROMPT_TEMPLATE
 
-        self.api.post(
+        if update:
+            method = self.api.put
+        else:
+            method = self.api.post
+
+        method(
             f'/projects/{self.project}/minds',
             data={
                 'name': name,

minds/rest_api.py

@@ -57,6 +57,16 @@ def post(self, url, data):
         _raise_for_status(resp)
         return resp
 
+    def put(self, url, data):
+        resp = requests.put(
+            self.base_url + url,
+            headers=self._headers(),
+            json=data,
+        )
+
+        _raise_for_status(resp)
+        return resp
+
     def patch(self, url, data):
         resp = requests.patch(
             self.base_url + url,

tests/integration/test_base_flow.py

@@ -37,7 +37,8 @@ def test_datasources():
 
     # create
     ds = client.datasources.create(example_ds)
-    ds = client.datasources.create(example_ds, replace=True)
+    assert ds.name == example_ds.name
+    ds = client.datasources.create(example_ds, update=True)
     assert ds.name == example_ds.name
 
     # get
@@ -90,6 +91,12 @@ def test_minds():
         datasources=[ds.name, ds2_cfg],
         prompt_template=prompt1
     )
+    mind = client.minds.create(
+        mind_name,
+        update=True,
+        datasources=[ds.name, ds2_cfg],
+        prompt_template=prompt1
+    )
 
     # get
     mind = client.minds.get(mind_name)

tests/unit/test_unit.py

@@ -36,9 +36,10 @@ def _compare_ds(self, ds1, ds2):
         assert ds1.tables == ds2.tables
 
     @patch('requests.get')
+    @patch('requests.put')
     @patch('requests.post')
     @patch('requests.delete')
-    def test_create_datasources(self, mock_del, mock_post, mock_get):
+    def test_create_datasources(self, mock_del, mock_post, mock_put, mock_get):
         client = get_client()
         response_mock(mock_get, example_ds.model_dump())
 
@@ -53,12 +54,9 @@ def check_ds_created(ds, mock_post):
 
         check_ds_created(ds, mock_post)
 
-        # with replace
-        ds = client.datasources.create(example_ds, replace=True)
-        args, _ = mock_del.call_args
-        assert args[0].endswith(f'/api/datasources/{example_ds.name}')
-
-        check_ds_created(ds, mock_post)
+        # with update
+        ds = client.datasources.create(example_ds, update=True)
+        check_ds_created(ds, mock_put)
 
     @patch('requests.get')
     def test_get_datasource(self, mock_get):
@@ -115,9 +113,10 @@ def compare_mind(self, mind, mind_json):
         assert mind.parameters == mind_json['parameters']
 
     @patch('requests.get')
+    @patch('requests.put')
     @patch('requests.post')
     @patch('requests.delete')
-    def test_create(self, mock_del, mock_post, mock_get):
+    def test_create(self, mock_del, mock_post, mock_put, mock_get):
         client = get_client()
 
         mind_name = 'test_mind'
@@ -145,7 +144,7 @@ def check_mind_created(mind, mock_post, create_params):
 
         check_mind_created(mind, mock_post, create_params)
 
-        # with replace
+        # -- with replace --
         create_params = {
             'name': mind_name,
             'prompt_template': prompt_template,
@@ -159,6 +158,14 @@ def check_mind_created(mind, mock_post, create_params):
 
         check_mind_created(mind, mock_post, create_params)
 
+        # -- with update --
+        mock_del.reset_mock()
+        mind = client.minds.create(update=True, **create_params)
+        # is not deleted
+        assert not mock_del.called
+
+        check_mind_created(mind, mock_put, create_params)
+
     @patch('requests.get')
     @patch('requests.patch')
     def test_update(self, mock_patch, mock_get):