upsert 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, replace=False, update=False):
         """
         Create new datasource and return it
 
@@ -30,6 +31,8 @@ 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 replace: if true - to remove existing datasource, default is false
+        :param update: if true - to update datasourse if exists, default is false
         :return: datasource object
         """
 
@@ -42,7 +45,10 @@ def create(self, ds_config: DatabaseConfig, replace=False):
             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/rest_api.py

@@ -53,6 +53,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,8 +37,11 @@ def test_datasources():
 
     # create
     ds = client.datasources.create(example_ds)
+    assert ds.name == example_ds.name
     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
     ds = client.datasources.get(example_ds.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())
 
@@ -60,6 +61,10 @@ def check_ds_created(ds, mock_post):
 
         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):
         client = get_client()