remove necessity of superuser to break the cache

Montana Low · Montana Low · commit 24427b2fb155 · 2022-08-22T10:40:39.000-07:00
diff --git a/pgml-extension/pgml_extension/model.py b/pgml-extension/pgml_extension/model.py
@@ -964,31 +964,27 @@ def deploy(self, strategy):
 
         plpy.execute(
             f"""
-        CREATE OR REPLACE FUNCTION pgml.predict(
-            project_name TEXT,          -- Human-friendly project name
-            features DOUBLE PRECISION[] -- Must match the training data column order
-        )
-        RETURNS DOUBLE PRECISION
-        AS $$
-            from pgml_extension.model import Project
-
-            return float(Project.find_by_name(project_name, {q(deployment_id)}).deployed_model.predict(features))
-        $$ LANGUAGE plpython3u;
+            CREATE OR REPLACE FUNCTION pgml.predict(
+                project_name TEXT,          -- Human-friendly project name
+                features DOUBLE PRECISION[] -- Must match the training data column order
+            ) 
+            RETURNS DOUBLE PRECISION
+            AS $$
+                SELECT pgml.predict_versioned(project_name, features, {q(deployment_id)});
+            $$ LANGUAGE SQL;
         """
         )
 
         plpy.execute(
             f"""
-        CREATE OR REPLACE FUNCTION pgml.predict_joint(
-            project_name TEXT,          -- Human-friendly project name
-            features DOUBLE PRECISION[] -- Must match the training data column order
-        )
-        RETURNS DOUBLE PRECISION[]
-        AS $$
-            from pgml_extension.model import Project
-
-            return Project.find_by_name(project_name, {q(deployment_id)}).deployed_model.predict(features)
-        $$ LANGUAGE plpython3u;
+            CREATE OR REPLACE FUNCTION pgml.predict_joint(
+                project_name TEXT,          -- Human-friendly project name
+                features DOUBLE PRECISION[] -- Must match the training data column order
+            ) 
+            RETURNS DOUBLE PRECISION[]
+            AS $$
+                SELECT pgml.predict_joint_versioned(project_name, features, {q(deployment_id)});
+            $$ LANGUAGE SQL;
         """
         )
 
diff --git a/pgml-extension/pgml_extension/sql/install/models.sql b/pgml-extension/pgml_extension/sql/install/models.sql
@@ -153,32 +153,50 @@ $$ LANGUAGE plpython3u;
 ---
 --- Predict
 ---
-CREATE OR REPLACE FUNCTION pgml.predict(
-	project_name TEXT,          -- Human-friendly project name
-	features DOUBLE PRECISION[] -- Must match the training data column order
+CREATE OR REPLACE FUNCTION pgml.predict_versioned(
+	project_name TEXT,           -- Human-friendly project name
+	features DOUBLE PRECISION[], -- Must match the training data column order
+	version INTEGER DEFAULT 0
 )
 RETURNS DOUBLE PRECISION
 AS $$
 	from pgml_extension.model import Project
 
-	return float(Project.find_by_name(project_name, 0).deployed_model.predict(features))
+	return float(Project.find_by_name(project_name, version).deployed_model.predict(features))
 $$ LANGUAGE plpython3u;
 
+CREATE OR REPLACE FUNCTION pgml.predict(
+	project_name TEXT,          -- Human-friendly project name
+	features DOUBLE PRECISION[] -- Must match the training data column order
+) 
+RETURNS DOUBLE PRECISION
+AS $$
+	SELECT pgml.predict_versioned(project_name, features, 0);
+$$ LANGUAGE SQL;
 
 ---
 --- Predict w/ multiple outputs
 ---
-CREATE OR REPLACE FUNCTION pgml.predict_joint(
-	project_name TEXT,          -- Human-friendly project name
-	features DOUBLE PRECISION[] -- Must match the training data column order
+CREATE OR REPLACE FUNCTION pgml.predict_joint_versioned(
+	project_name TEXT,           -- Human-friendly project name
+	features DOUBLE PRECISION[], -- Must match the training data column order
+	version INTEGER DEFAULT 0
 )
 RETURNS DOUBLE PRECISION[]
 AS $$
 	from pgml_extension.model import Project
 
-	return Project.find_by_name(project_name, 0).deployed_model.predict(features)
+	return Project.find_by_name(project_name, version).deployed_model.predict(features)
 $$ LANGUAGE plpython3u;
 
+CREATE OR REPLACE FUNCTION pgml.predict_joint(
+	project_name TEXT,          -- Human-friendly project name
+	features DOUBLE PRECISION[] -- Must match the training data column order
+) 
+RETURNS DOUBLE PRECISION[]
+AS $$
+	SELECT pgml.predict_joint_versioned(project_name, features, 0);
+$$ LANGUAGE SQL;
 
 ---
 --- Predict using a specific model. Useful for debugging.
