more fixes

levkk · levkk · commit 28545c9b39a7 · 2022-04-14T15:10:09.000-07:00
diff --git a/pgml/pgml/model.py b/pgml/pgml/model.py
@@ -208,11 +208,11 @@ def data(self):
 
         # Sanity check the data
         if len(data) == 0:
-            PgMLException(
+            raise PgMLException(
                 f"Relation `{self.relation_name}` contains no rows. Did you pass the correct `relation_name`?"
             )
         if self.y_column_name not in data[0]:
-            PgMLException(
+            raise PgMLException(
                 f"Column `{self.y_column_name}` not found. Did you pass the correct `y_column_name`?"
             )
 
@@ -428,6 +428,10 @@ def train(
         algorithms = ["linear", "random_forest"]
     elif objective == "classification":
         algorithms = ["random_forest"]
+    else:
+        raise PgMLException(
+            f"Unknown objective '{objective}', available options are: regression, classification"
+        )
 
     for algorithm_name in algorithms:
         model = Model.create(project, snapshot, algorithm_name)
diff --git a/sql/install.sql b/sql/install.sql
@@ -103,15 +103,15 @@ $$ LANGUAGE plpython3u;
 ---
 --- Regression
 ---
-DROP FUNCTION pgml.train(project_name TEXT, objective TEXT, relation_name TEXT, y_column_name TEXT);
+DROP FUNCTION IF EXISTS pgml.train(project_name TEXT, objective TEXT, relation_name TEXT, y_column_name TEXT);
 CREATE OR REPLACE FUNCTION pgml.train(project_name TEXT, objective TEXT, relation_name TEXT, y_column_name TEXT)
-RETURNS TEXT
+RETURNS TABLE(project_name TEXT, objective TEXT, status TEXT)
 AS $$
 	from pgml.model import train
 
 	train(project_name, objective, relation_name, y_column_name)
 
-	return "OK"
+	return [(project_name, objective, "deployed")]
 $$ LANGUAGE plpython3u;
 
 ---
