12912993071

Committed 22 Jan 2025 04:46PM UTC coverage: 90.038% (-0.02%) from 90.061%

Build # 12912993071

Build Type

Pull #1007

github

Committed by

web-flow

Commit Message

Merge 7955102e5 into df8c694c8

Pull Request Pull Request #1007: Add ml permissions

Run Details

30 of 71 new or added lines in 5 files covered. (42.25%)

71 existing lines in 6 files now uncovered.

127254 of 141333 relevant lines covered (90.04%)

56955.06 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

75.74

/services/src/machine_learning/postgres.rs

use crate::{
    machine_learning::{
        error::{
            error::{Bb8MachineLearningError, PostgresMachineLearningError},
            MachineLearningError,
        },
        name::MlModelName,
        MlModel, MlModelDb, MlModelId, MlModelIdAndName, MlModelListOptions, MlModelMetadata,
    },
    permissions::Permission,
    pro::contexts::PostgresDb,
    util::postgres::PostgresErrorExt,
};
use async_trait::async_trait;
use geoengine_datatypes::util::Identifier;
use snafu::ResultExt;
use tokio_postgres::{
    tls::{MakeTlsConnect, TlsConnect},
    Socket,
};

#[async_trait]
impl<Tls> MlModelDb for PostgresDb<Tls>
where
    Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static + std::fmt::Debug,
    <Tls as MakeTlsConnect<Socket>>::Stream: Send + Sync,
    <Tls as MakeTlsConnect<Socket>>::TlsConnect: Send,
    <<Tls as MakeTlsConnect<Socket>>::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
    async fn list_models(
        &self,
        options: &MlModelListOptions,
    ) -> Result<Vec<MlModel>, MachineLearningError> {
        let conn = self
            .conn_pool
            .get()
            .await
            .context(Bb8MachineLearningError)?;

        let rows = conn
            .query(
                "
                SELECT
                    m.id,
                    m.name,
                    m.display_name,
                    m.description,
                    m.upload,
                    m.metadata
                FROM 
                    user_permitted_ml_models u JOIN ml_models m ON (u.ml_model_id = m.id)
                WHERE 
                    u.user_id = $1
                ORDER BY
                    name
                OFFSET
                    $2
                LIMIT 
                    $3",
                &[
                    &self.session.user.id,
                    &i64::from(options.offset),
                    &i64::from(options.limit),
                ],
            )
            .await
            .context(PostgresMachineLearningError)?;

        let models = rows
            .into_iter()
            .map(|row| MlModel {
                name: row.get(1),
                display_name: row.get(2),
                description: row.get(3),
                upload: row.get(4),
                metadata: row.get(5),
            })
            .collect();

        Ok(models)
    }

    async fn load_model(&self, name: &MlModelName) -> Result<MlModel, MachineLearningError> {
        let conn = self
            .conn_pool
            .get()
            .await
            .context(Bb8MachineLearningError)?;

        let Some(row) = conn
            .query_opt(
                "SELECT
                    m.id,
                    m.name,
                    m.display_name,
                    m.description,
                    m.upload,
                    m.metadata
                FROM 
                    user_permitted_ml_models u JOIN ml_models m ON (u.ml_model_id = m.id)
                WHERE 
                    u.user_id = $1 AND m.name = $2::\"MlModelName\"",
                &[&self.session.user.id, name],
            )
            .await
            .context(PostgresMachineLearningError)?
        else {
            return Err(MachineLearningError::ModelNotFound { name: name.clone() });
        };

        Ok(MlModel {
            name: row.get(1),
            display_name: row.get(2),
            description: row.get(3),
            upload: row.get(4),
            metadata: row.get(5),
        })
    }

    async fn load_model_metadata(
        &self,
        name: &MlModelName,
    ) -> Result<MlModelMetadata, MachineLearningError> {
        let conn = self
            .conn_pool
            .get()
            .await
            .context(Bb8MachineLearningError)?;

        let Some(row) = conn
            .query_opt(
                "SELECT
                    m.metadata
                FROM 
                    user_permitted_ml_models u JOIN ml_models m ON (u.ml_model_id = m.id)
                WHERE 
                    u.user_id = $1 AND m.name = $2",
                &[&self.session.user.id, name],
            )
            .await
            .context(PostgresMachineLearningError)?
        else {
            return Err(MachineLearningError::ModelNotFound { name: name.clone() });
        };

        Ok(row.get(1))
    }

    async fn add_model(&self, model: MlModel) -> Result<MlModelIdAndName, MachineLearningError> {
        self.check_ml_model_namespace(&model.name)?;

        let mut conn = self
            .conn_pool
            .get()
            .await
            .context(Bb8MachineLearningError)?;

        let tx = conn
            .build_transaction()
            .start()
            .await
            .context(PostgresMachineLearningError)?;

        let id = MlModelId::new();

        tx.query_one(
            "INSERT INTO ml_models (
                    id,
                    name, 
                    display_name, 
                    description,
                    upload,
                    metadata
                ) VALUES ($1, $2, $3, $4, $5, $6) RETURNING id;",
            &[
                &id,
                &model.name,
                &model.display_name,
                &model.description,
                &model.upload,
                &model.metadata,
            ],
        )
        .await
        .map_unique_violation("ml_models", "name", || {
            MachineLearningError::DuplicateMlModelName {
                name: model.name.clone(),
            }
        })?;

        let stmt = tx
            .prepare(
                "INSERT INTO permissions (role_id, permission, ml_model_id) VALUES ($1, $2, $3);",
            )
            .await
            .context(PostgresMachineLearningError)?;

        tx.execute(&stmt, &[&self.session.user.id, &Permission::Owner, &id])
            .await
            .context(PostgresMachineLearningError)?;

        tx.commit().await.context(PostgresMachineLearningError)?;

        Ok(MlModelIdAndName {
            id,
            name: model.name,
        })
    }

    async fn resolve_model_name_to_id(
        &self,
        model_name: &MlModelName,
    ) -> Result<Option<MlModelId>, MachineLearningError> {
        let conn = self
            .conn_pool
            .get()
            .await
            .context(Bb8MachineLearningError)?;

        let stmt = conn
            .prepare(
                "SELECT id
        FROM ml_models
        WHERE name = $1::\"MlModelName\"",
            )
            .await?;

        let row_option = conn.query_opt(&stmt, &[&model_name]).await?;

        Ok(row_option.map(|row| row.get(0)))
    }
}

1	use crate::{
2	machine_learning::{
3	error::{
4	error::{Bb8MachineLearningError, PostgresMachineLearningError},
5	MachineLearningError,
6	},
7	name::MlModelName,
8	MlModel, MlModelDb, MlModelId, MlModelIdAndName, MlModelListOptions, MlModelMetadata,
9	},
10	permissions::Permission,
11	pro::contexts::PostgresDb,
12	util::postgres::PostgresErrorExt,
13	};
14	use async_trait::async_trait;
15	use geoengine_datatypes::util::Identifier;
16	use snafu::ResultExt;
17	use tokio_postgres::{
18	tls::{MakeTlsConnect, TlsConnect},
19	Socket,
20	};
21
22	#[async_trait]
23	impl<Tls> MlModelDb for PostgresDb<Tls>
24	where
25	Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static + std::fmt::Debug,
26	<Tls as MakeTlsConnect<Socket>>::Stream: Send + Sync,
27	<Tls as MakeTlsConnect<Socket>>::TlsConnect: Send,
28	<<Tls as MakeTlsConnect<Socket>>::TlsConnect as TlsConnect<Socket>>::Future: Send,
29	{
30	async fn list_models(
31	&self,
32	options: &MlModelListOptions,
33	) -> Result<Vec<MlModel>, MachineLearningError> {	1✔
34	let conn = self	1✔
35	.conn_pool	1✔
36	.get()	1✔
37	.await	1✔
38	.context(Bb8MachineLearningError)?;	1✔
39
40	let rows = conn	1✔
41	.query(	1✔
42	"	1✔
43	SELECT	1✔
44	m.id,	1✔
45	m.name,	1✔
46	m.display_name,	1✔
47	m.description,	1✔
48	m.upload,	1✔
49	m.metadata	1✔
50	FROM	1✔
51	user_permitted_ml_models u JOIN ml_models m ON (u.ml_model_id = m.id)	1✔
52	WHERE	1✔
53	u.user_id = $1	1✔
54	ORDER BY	1✔
55	name	1✔
56	OFFSET	1✔
57	$2	1✔
58	LIMIT	1✔
59	$3",	1✔
60	&[	1✔
61	&self.session.user.id,	1✔
62	&i64::from(options.offset),	1✔
63	&i64::from(options.limit),	1✔
64	],	1✔
65	)	1✔
66	.await	1✔
67	.context(PostgresMachineLearningError)?;	1✔
68
69	let models = rows	1✔
70	.into_iter()	1✔
71	.map(\|row\| MlModel {	1✔
72	name: row.get(1),	1✔
73	display_name: row.get(2),	1✔
74	description: row.get(3),	1✔
75	upload: row.get(4),	1✔
76	metadata: row.get(5),	1✔
77	})	1✔
78	.collect();	1✔
79		1✔
80	Ok(models)	1✔
81	}	2✔
82
83	async fn load_model(&self, name: &MlModelName) -> Result<MlModel, MachineLearningError> {	1✔
84	let conn = self	1✔
85	.conn_pool	1✔
86	.get()	1✔
87	.await	1✔
88	.context(Bb8MachineLearningError)?;	1✔
89
90	let Some(row) = conn	1✔
91	.query_opt(	1✔
92	"SELECT	1✔
93	m.id,	1✔
94	m.name,	1✔
95	m.display_name,	1✔
96	m.description,	1✔
97	m.upload,	1✔
98	m.metadata	1✔
99	FROM	1✔
100	user_permitted_ml_models u JOIN ml_models m ON (u.ml_model_id = m.id)	1✔
101	WHERE	1✔
102	u.user_id = $1 AND m.name = $2::\"MlModelName\"",	1✔
103	&[&self.session.user.id, name],	1✔
104	)	1✔
105	.await	1✔
106	.context(PostgresMachineLearningError)?	1✔
107	else {
108	return Err(MachineLearningError::ModelNotFound { name: name.clone() });	×
109	};
110
111	Ok(MlModel {	1✔
112	name: row.get(1),	1✔
113	display_name: row.get(2),	1✔
114	description: row.get(3),	1✔
115	upload: row.get(4),	1✔
116	metadata: row.get(5),	1✔
117	})	1✔
118	}	2✔
119
120	async fn load_model_metadata(
121	&self,
122	name: &MlModelName,
123	) -> Result<MlModelMetadata, MachineLearningError> {	×
124	let conn = self	×
125	.conn_pool	×
126	.get()	×
127	.await	×
128	.context(Bb8MachineLearningError)?;	×
129
130	let Some(row) = conn	×
131	.query_opt(	×
132	"SELECT	×
133	m.metadata	×
134	FROM	×
135	user_permitted_ml_models u JOIN ml_models m ON (u.ml_model_id = m.id)	×
136	WHERE	×
137	u.user_id = $1 AND m.name = $2",	×
138	&[&self.session.user.id, name],	×
139	)	×
140	.await	×
141	.context(PostgresMachineLearningError)?	×
142	else {
143	return Err(MachineLearningError::ModelNotFound { name: name.clone() });	×
144	};
145
146	Ok(row.get(1))	×
147	}	×
148
149	async fn add_model(&self, model: MlModel) -> Result<MlModelIdAndName, MachineLearningError> {	1✔
150	self.check_ml_model_namespace(&model.name)?;	1✔
151
152	let mut conn = self	1✔
153	.conn_pool	1✔
154	.get()	1✔
155	.await	1✔
156	.context(Bb8MachineLearningError)?;	1✔
157
158	let tx = conn	1✔
159	.build_transaction()	1✔
160	.start()	1✔
161	.await	1✔
162	.context(PostgresMachineLearningError)?;	1✔
163
164	let id = MlModelId::new();	1✔
165		1✔
166	tx.query_one(	1✔
167	"INSERT INTO ml_models (	1✔
168	id,	1✔
169	name,	1✔
170	display_name,	1✔
171	description,	1✔
172	upload,	1✔
173	metadata	1✔
174	) VALUES ($1, $2, $3, $4, $5, $6) RETURNING id;",	1✔
175	&[	1✔
176	&id,	1✔
177	&model.name,	1✔
178	&model.display_name,	1✔
179	&model.description,	1✔
180	&model.upload,	1✔
181	&model.metadata,	1✔
182	],	1✔
183	)	1✔
184	.await	1✔
185	.map_unique_violation("ml_models", "name", \|\| {	1✔
186	MachineLearningError::DuplicateMlModelName {	×
187	name: model.name.clone(),	×
188	}	×
189	})?;	1✔
190
191	let stmt = tx	1✔
192	.prepare(	1✔
193	"INSERT INTO permissions (role_id, permission, ml_model_id) VALUES ($1, $2, $3);",	1✔
194	)	1✔
195	.await	1✔
196	.context(PostgresMachineLearningError)?;	1✔
197
198	tx.execute(&stmt, &[&self.session.user.id, &Permission::Owner, &id])	1✔
199	.await	1✔
200	.context(PostgresMachineLearningError)?;	1✔
201
202	tx.commit().await.context(PostgresMachineLearningError)?;	1✔
203
204	Ok(MlModelIdAndName {	1✔
205	id,	1✔
206	name: model.name,	1✔
207	})	1✔
208	}	2✔
209
210	async fn resolve_model_name_to_id(
211	&self,
212	model_name: &MlModelName,
NEW 213	) -> Result<Option<MlModelId>, MachineLearningError> {	×
NEW 214	let conn = self	×
NEW 215	.conn_pool	×
NEW 216	.get()	×
NEW 217	.await	×
NEW 218	.context(Bb8MachineLearningError)?;	×
219
NEW 220	let stmt = conn	×
NEW 221	.prepare(	×
NEW 222	"SELECT id	×
NEW 223	FROM ml_models	×
NEW 224	WHERE name = $1::\"MlModelName\"",	×
NEW 225	)	×
NEW 226	.await?;	×
227
NEW 228	let row_option = conn.query_opt(&stmt, &[&model_name]).await?;	×
229
NEW 230	Ok(row_option.map(\|row\| row.get(0)))	×
UNCOV 231	}	×
232	}

geo-engine / geoengine / 12912993071

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous