Coverage for pepdbagent/modules/sample.py: 92%
90 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-01 22:07 +0000
1import datetime
2import logging
4import peprs
5from sqlalchemy import and_, select
6from sqlalchemy.orm import Session
7from sqlalchemy.orm.attributes import flag_modified
9from pepdbagent.const import DEFAULT_TAG, PKG_NAME, SAMPLE_TABLE_INDEX_KEY
10from pepdbagent.db_utils import BaseEngine, Projects, Samples
11from pepdbagent.exceptions import SampleAlreadyExistsError, SampleNotFoundError
12from pepdbagent.utils import generate_guid, order_samples
14_LOGGER = logging.getLogger(PKG_NAME)
17class PEPDatabaseSample:
18 """
19 Class that represents Project in Database.
21 While using this class, user can create, retrieve, delete, and update projects from database
22 """
24 def __init__(self, pep_db_engine: BaseEngine):
25 """
26 Args:
27 pep_db_engine: PEPDatabaseAgent engine object.
28 """
29 self._sa_engine = pep_db_engine.engine
30 self._pep_db_engine = pep_db_engine
32 def get(
33 self,
34 namespace: str,
35 name: str,
36 sample_name: str,
37 tag: str = DEFAULT_TAG,
38 raw: bool = True,
39 ) -> peprs.Sample | dict | None:
40 """Retrieve a sample from the database.
42 Args:
43 namespace: Namespace of the project.
44 name: Name of the project.
45 sample_name: Name of the sample.
46 tag: Tag of the project.
47 raw: Return raw dict if True, peprs.Sample object if False.
49 Returns:
50 Raw sample dict or peprs.Sample object.
52 Raises:
53 SampleNotFoundError: If the sample does not exist.
54 """
55 statement_sample = select(Samples).where(
56 and_(
57 Samples.project_id
58 == select(Projects.id)
59 .where(
60 and_(
61 Projects.namespace == namespace,
62 Projects.name == name,
63 Projects.tag == tag,
64 ),
65 )
66 .scalar_subquery(),
67 Samples.sample_name == sample_name,
68 )
69 )
70 project_config_statement = select(Projects.config).where(
71 and_(
72 Projects.namespace == namespace,
73 Projects.name == name,
74 Projects.tag == tag,
75 )
76 )
78 with Session(self._sa_engine) as session:
79 result = session.scalar(statement_sample)
80 if result:
81 if not raw:
82 config = session.execute(project_config_statement).one_or_none()[0]
83 project = peprs.Project.from_dict(
84 pep_dictionary={
85 "config": config,
86 "samples": [result.sample],
87 }
88 )
89 return project.samples[0]
90 else:
91 return result.sample
92 else:
93 raise SampleNotFoundError(
94 f"Sample {namespace}/{name}:{tag}?{sample_name} not found in the database"
95 )
97 def update(
98 self,
99 namespace: str,
100 name: str,
101 tag: str,
102 sample_name: str,
103 update_dict: dict,
104 full_update: bool = False,
105 ) -> None:
106 """Update a sample in the database.
108 Args:
109 namespace: Namespace of the project.
110 name: Name of the project.
111 tag: Tag of the project.
112 sample_name: Name of the sample.
113 update_dict: Dict of fields to update, e.g., {"sample_name": "s1", "protocol": "rna"}.
114 full_update: Replace all sample fields if True, merge if False.
116 Raises:
117 SampleNotFoundError: If the sample does not exist.
118 """
119 statement = select(Samples).where(
120 and_(
121 Samples.project_id
122 == select(Projects.id)
123 .where(
124 and_(
125 Projects.namespace == namespace,
126 Projects.name == name,
127 Projects.tag == tag,
128 ),
129 )
130 .scalar_subquery(),
131 Samples.sample_name == sample_name,
132 )
133 )
134 project_statement = select(Projects).where(
135 and_(
136 Projects.namespace == namespace,
137 Projects.name == name,
138 Projects.tag == tag,
139 )
140 )
141 with Session(self._sa_engine) as session:
142 sample_mapping = session.scalar(statement)
143 project_mapping = session.scalar(project_statement)
145 if sample_mapping:
146 if full_update:
147 sample_mapping.sample = update_dict
148 else:
149 sample_mapping.sample.update(update_dict)
150 try:
151 sample_mapping.sample_name = sample_mapping.sample[
152 project_mapping.config.get(
153 SAMPLE_TABLE_INDEX_KEY, "sample_name"
154 )
155 ]
156 except KeyError:
157 raise KeyError(
158 f"Sample index key {project_mapping.config.get('sample_table_index', 'sample_name')} not found in sample dict"
159 )
161 # This line needed due to: https://github.com/sqlalchemy/sqlalchemy/issues/5218
162 flag_modified(sample_mapping, "sample")
164 project_mapping.last_update_date = datetime.datetime.now(
165 datetime.timezone.utc
166 )
168 session.commit()
169 else:
170 raise SampleNotFoundError(
171 f"Sample {namespace}/{name}:{tag}?{sample_name} not found in the database"
172 )
174 def add(
175 self,
176 namespace: str,
177 name: str,
178 tag: str,
179 sample_dict: dict,
180 overwrite: bool = False,
181 ) -> None:
182 """Add a sample to a project in the database.
184 Args:
185 namespace: Namespace of the project.
186 name: Name of the project.
187 tag: Tag of the project.
188 sample_dict: Sample data, e.g., {"sample_name": "s1", "protocol": "rna"}.
189 overwrite: Overwrite the sample if it already exists.
191 Raises:
192 SampleAlreadyExistsError: If the sample exists and overwrite is False.
193 """
195 with Session(self._sa_engine) as session:
196 project_statement = select(Projects).where(
197 and_(
198 Projects.namespace == namespace,
199 Projects.name == name,
200 Projects.tag == tag,
201 )
202 )
203 # project mapping is needed to update number of samples, last_update_date and get sample_index_key
204 project_mapping = session.scalar(project_statement)
205 try:
206 sample_name = sample_dict[
207 project_mapping.config.get(SAMPLE_TABLE_INDEX_KEY, "sample_name")
208 ]
209 except KeyError:
210 raise KeyError(
211 f"Sample index key {project_mapping.config.get('sample_table_index', 'sample_name')} not found in sample dict"
212 )
213 statement = select(Samples).where(
214 and_(
215 Samples.project_id == project_mapping.id,
216 Samples.sample_name == sample_name,
217 )
218 )
219 sample_mapping = session.scalar(statement)
221 if sample_mapping and not overwrite:
222 raise SampleAlreadyExistsError(
223 f"Sample {namespace}/{name}:{tag}?{sample_name} already exists in the database"
224 )
225 elif sample_mapping and overwrite:
226 self.update(
227 namespace=namespace,
228 name=name,
229 tag=tag,
230 sample_name=sample_name,
231 update_dict=sample_dict,
232 full_update=True,
233 )
234 return None
235 else:
236 sample_mapping = Samples(
237 sample=sample_dict,
238 project_id=project_mapping.id,
239 sample_name=sample_name,
240 guid=generate_guid(),
241 parent_guid=self._get_last_sample_guid(project_mapping.id),
242 )
243 project_mapping.number_of_samples += 1
244 project_mapping.last_update_date = datetime.datetime.now(
245 datetime.timezone.utc
246 )
248 session.add(sample_mapping)
249 session.commit()
251 def _get_last_sample_guid(self, project_id: int) -> str:
252 """Get the guid of the last sample in the project chain.
254 Args:
255 project_id: Database id of the project.
257 Returns:
258 GUID of the last sample.
259 """
260 statement = select(Samples).where(Samples.project_id == project_id)
261 with Session(self._sa_engine) as session:
262 samples_results = session.scalars(statement)
264 result_dict = {}
265 for sample in samples_results:
266 sample_dict = sample.sample
268 result_dict[sample.guid] = {
269 "sample": sample_dict,
270 "guid": sample.guid,
271 "parent_guid": sample.parent_guid,
272 }
273 return order_samples(result_dict)[-1]["guid"]
275 def delete(
276 self,
277 namespace: str,
278 name: str,
279 tag: str,
280 sample_name: str,
281 ) -> None:
282 """Delete a sample from the database.
284 Args:
285 namespace: Namespace of the project.
286 name: Name of the project.
287 tag: Tag of the project.
288 sample_name: Name of the sample.
290 Raises:
291 SampleNotFoundError: If the sample does not exist.
292 """
293 statement = select(Samples).where(
294 and_(
295 Samples.project_id
296 == select(Projects.id)
297 .where(
298 and_(
299 Projects.namespace == namespace,
300 Projects.name == name,
301 Projects.tag == tag,
302 ),
303 )
304 .scalar_subquery(),
305 Samples.sample_name == sample_name,
306 )
307 )
308 project_statement = select(Projects).where(
309 and_(
310 Projects.namespace == namespace,
311 Projects.name == name,
312 Projects.tag == tag,
313 )
314 )
315 with Session(self._sa_engine) as session:
316 sample_mapping = session.scalar(statement)
317 project_mapping = session.scalar(project_statement)
319 if sample_mapping:
320 parent_mapping = sample_mapping.parent_mapping
321 child_mapping = sample_mapping.child_mapping
322 session.delete(sample_mapping)
323 if child_mapping:
324 child_mapping.parent_mapping = parent_mapping
325 project_mapping.number_of_samples -= 1
326 project_mapping.last_update_date = datetime.datetime.now(
327 datetime.timezone.utc
328 )
329 session.commit()
330 else:
331 raise SampleNotFoundError(
332 f"Sample {namespace}/{name}:{tag}?{sample_name} not found in the database"
333 )