Merge pull request #7 from mindsdb/exceptions

Fix: add_datasource to mind
mindsdb · Sep 25, 2024 · e3fa8c7 · e3fa8c7
2 parents 0922300 + 295c33c
commit e3fa8c7
Showing 1 changed file with 24 additions and 25 deletions.
diff --git a/minds/minds.py b/minds/minds.py
@@ -47,18 +47,7 @@ def update(
         if datasources is not None:
             ds_names = []
             for ds in datasources:
-                if isinstance(ds, Datasource):
-                    ds = ds.name
-                elif isinstance(ds, DatabaseConfig):
-                    # if not exists - create
-                    try:
-                        self.client.datasources.get(ds.name)
-                    except exc.ObjectNotFound:
-                        self.client.datasources.create(ds)
-
-                    ds = ds.name
-                elif not isinstance(ds, str):
-                    raise ValueError(f'Unknown type of datasource: {ds}')
+                ds = self.client.minds._check_datasource(ds)
                 ds_names.append(ds)
             data['datasources'] = ds_names
 
@@ -79,9 +68,14 @@ def update(
             self.name = name
 
     def add_datasource(self, datasource: Datasource):
+
+        ds_name = self.client.minds._check_datasource(datasource)
+
         self.api.post(
             f'/projects/{self.project}/minds/{self.name}/datasources',
-            data=datasource.model_dump()
+            data={
+                'name': ds_name,
+            }
         )
         updated = self.client.minds.get(self.name)
 
@@ -135,6 +129,21 @@ def get(self, name: str) -> Mind:
         item = self.api.get(f'/projects/{self.project}/minds/{name}').json()
         return Mind(self.client, **item)
 
+    def _check_datasource(self, ds) -> str:
+        if isinstance(ds, Datasource):
+            ds = ds.name
+        elif isinstance(ds, DatabaseConfig):
+            # if not exists - create
+            try:
+                self.client.datasources.get(ds.name)
+            except exc.ObjectNotFound:
+                self.client.datasources.create(ds)
+
+            ds = ds.name
+        elif not isinstance(ds, str):
+            raise ValueError(f'Unknown type of datasource: {ds}')
+        return ds
+
     def create(
         self, name,
         model_name=None,
@@ -154,18 +163,8 @@ def create(
         ds_names = []
         if datasources:
             for ds in datasources:
-                if isinstance(ds, Datasource):
-                    ds = ds.name
-                elif isinstance(ds, DatabaseConfig):
-                    # if not exists - create
-                    try:
-                        self.client.datasources.get(ds.name)
-                    except exc.ObjectNotFound:
-                        self.client.datasources.create(ds)
-
-                    ds = ds.name
-                elif not isinstance(ds, str):
-                    raise ValueError(f'Unknown type of datasource: {ds}')
+                ds = self._check_datasource(ds)
+
                 ds_names.append(ds)
 
         if parameters is None: