Coverage for pepdbagent/modules/annotation.py: 75%
186 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-01 22:07 +0000
1import logging
2from datetime import datetime
3from typing import Literal
5from sqlalchemy import and_, func, or_, select
6from sqlalchemy.orm import Session
7from sqlalchemy.sql.selectable import Select
9from pepdbagent.const import (
10 DEFAULT_LIMIT,
11 DEFAULT_OFFSET,
12 DEFAULT_TAG,
13 LAST_UPDATE_DATE_KEY,
14 PKG_NAME,
15 SUBMISSION_DATE_KEY,
16)
17from pepdbagent.db_utils import BaseEngine, Projects
18from pepdbagent.exceptions import FilterError, ProjectNotFoundError, RegistryPathError
19from pepdbagent.models import AnnotationList, AnnotationModel, RegistryPath
20from pepdbagent.utils import (
21 convert_date_string_to_date,
22 registry_path_converter,
23 tuple_converter,
24)
26_LOGGER = logging.getLogger(PKG_NAME)
29class PEPDatabaseAnnotation:
30 """
31 Class that represents project Annotations in the Database.
33 While using this class, user can retrieve all necessary metadata about PEPs
34 """
36 def __init__(self, pep_db_engine: BaseEngine):
37 """
38 Args:
39 pep_db_engine: PEPDatabaseAgent engine object.
40 """
41 self._sa_engine = pep_db_engine.engine
42 self._pep_db_engine = pep_db_engine
44 def get(
45 self,
46 namespace: str = None,
47 name: str = None,
48 tag: str = None,
49 query: str = None,
50 admin: list[str] | str | None = None,
51 limit: int = DEFAULT_LIMIT,
52 offset: int = DEFAULT_OFFSET,
53 order_by: str = "update_date",
54 order_desc: bool = False,
55 filter_by: Literal["submission_date", "last_update_date"] | None = None,
56 filter_start_date: str | None = None,
57 filter_end_date: str | None = None,
58 pep_type: Literal["pep", "pop"] | None = None,
59 ) -> AnnotationList:
60 """Get project annotations.
62 Five retrieval scenarios:
63 - namespace + name + tag: exact match.
64 - namespace only: all projects in that namespace.
65 - nothing: all projects in the database.
66 - query: full-text search across name, tag, description.
67 - query + namespace: full-text search within a namespace.
69 Args:
70 namespace: Namespace to filter by.
71 name: Project name (use with namespace and tag for exact lookup).
72 tag: Project tag.
73 query: Search string matched against name, tag, and description.
74 admin: Namespace(s) where the caller has admin rights.
75 limit: Maximum number of results.
76 offset: Number of results to skip.
77 order_by: Sort field — "name", "update_date", or "submission_date".
78 order_desc: Sort in descending order if True.
79 filter_by: Date field to apply range filter on — "submission_date" or "last_update_date".
80 filter_start_date: Range start in YYYY/MM/DD format.
81 filter_end_date: Range end in YYYY/MM/DD format (default: today).
82 pep_type: Restrict to "pep" or "pop" (default: all).
84 Returns:
85 AnnotationList with count, limit, offset, and results.
86 """
87 if all([namespace, name, tag]):
88 found_annotation = [
89 self._get_single_annotation(
90 namespace=namespace,
91 name=name,
92 tag=tag,
93 admin=admin,
94 )
95 ]
96 return AnnotationList(
97 count=len(found_annotation),
98 limit=1,
99 offset=0,
100 results=found_annotation,
101 )
103 if pep_type not in [None, "pep", "pop"]:
104 raise ValueError(
105 f"pep_type should be one of ['pep', 'pop'], got {pep_type}"
106 )
108 return AnnotationList(
109 limit=limit,
110 offset=offset,
111 count=self._count_projects(
112 namespace=namespace,
113 search_str=query,
114 tag=tag,
115 admin=admin,
116 filter_by=filter_by,
117 filter_end_date=filter_end_date,
118 filter_start_date=filter_start_date,
119 pep_type=pep_type,
120 ),
121 results=self._get_projects(
122 namespace=namespace,
123 search_str=query,
124 tag=tag,
125 admin=admin,
126 offset=offset,
127 limit=limit,
128 order_by=order_by,
129 order_desc=order_desc,
130 filter_by=filter_by,
131 filter_end_date=filter_end_date,
132 filter_start_date=filter_start_date,
133 pep_type=pep_type,
134 ),
135 )
137 def get_by_rp(
138 self,
139 registry_paths: list[str] | str,
140 admin: list[str] | str | None = None,
141 ) -> AnnotationList:
142 """Get project annotations by registry path or list of registry paths.
144 Args:
145 registry_paths: Single registry path or list of paths ("namespace/name:tag").
146 admin: Namespace(s) where the caller has admin rights.
148 Returns:
149 AnnotationList with count, limit, offset, and results.
150 """
151 if isinstance(registry_paths, list):
152 anno_results = []
153 for path in registry_paths:
154 try:
155 namespace, name, tag = registry_path_converter(path)
156 except RegistryPathError as err:
157 _LOGGER.error(str(err), registry_paths)
158 continue
159 try:
160 single_return = self._get_single_annotation(
161 namespace, name, tag, admin
162 )
163 if single_return:
164 anno_results.append(single_return)
165 except ProjectNotFoundError:
166 pass
167 return_len = len(anno_results)
168 return AnnotationList(
169 count=return_len,
170 limit=len(registry_paths),
171 offset=0,
172 results=anno_results,
173 )
175 else:
176 namespace, name, tag = registry_path_converter(registry_paths)
177 return self.get(namespace=namespace, name=name, tag=tag, admin=admin)
179 def _get_single_annotation(
180 self,
181 namespace: str,
182 name: str,
183 tag: str = DEFAULT_TAG,
184 admin: list[str] | str | None = None,
185 ) -> AnnotationModel | None:
186 """Retrieve annotation for a single project.
188 Args:
189 namespace: Project namespace.
190 name: Project name.
191 tag: Project tag.
192 admin: Namespace(s) with admin/private access.
194 Returns:
195 AnnotationModel for the project, or None if not found.
197 Raises:
198 ProjectNotFoundError: If the project does not exist or is not accessible.
199 """
200 _LOGGER.info(f"Getting annotation of the project: '{namespace}/{name}:{tag}'")
201 admin_tuple = tuple_converter(admin)
203 statement = select(Projects).where(
204 and_(
205 Projects.name == name,
206 Projects.namespace == namespace,
207 Projects.tag == tag,
208 or_(
209 Projects.namespace.in_(admin_tuple),
210 Projects.private.is_(False),
211 ),
212 )
213 )
214 with Session(self._sa_engine) as session:
215 query_result = session.scalar(statement)
217 if query_result:
218 annot = AnnotationModel(
219 namespace=query_result.namespace,
220 name=query_result.name,
221 tag=query_result.tag,
222 is_private=query_result.private,
223 description=query_result.description,
224 number_of_samples=query_result.number_of_samples,
225 submission_date=str(query_result.submission_date),
226 last_update_date=str(query_result.last_update_date),
227 digest=query_result.digest,
228 pep_schema=(
229 f"{query_result.schema_mapping.schema_mapping.namespace}/{query_result.schema_mapping.schema_mapping.name}:{query_result.schema_mapping.version}"
230 if query_result.schema_mapping
231 else None
232 ),
233 pop=query_result.pop,
234 stars_number=query_result.number_of_stars,
235 forked_from=(
236 f"{query_result.forked_from_mapping.namespace}/{query_result.forked_from_mapping.name}:{query_result.forked_from_mapping.tag}"
237 if query_result.forked_from_id
238 else None
239 ),
240 )
241 _LOGGER.info(
242 f"Annotation of the project '{namespace}/{name}:{tag}' has been found!"
243 )
244 return annot
245 else:
246 raise ProjectNotFoundError(
247 f"Project '{namespace}/{name}:{tag}' was not found."
248 )
250 def _count_projects(
251 self,
252 namespace: str = None,
253 search_str: str = None,
254 tag: str = None,
255 admin: str | list[str] | None = None,
256 filter_by: Literal["submission_date", "last_update_date"] | None = None,
257 filter_start_date: str | None = None,
258 filter_end_date: str | None = None,
259 pep_type: Literal["pep", "pop"] | None = None,
260 ) -> int:
261 """Count projects matching the given filters.
263 Args:
264 namespace: Namespace to restrict the search to.
265 search_str: String to search in name, tag, and description.
266 tag: Exact tag to match.
267 admin: Namespace(s) with admin/private access.
268 filter_by: Date field for range filter — "submission_date" or "last_update_date".
269 filter_start_date: Range start in YYYY/MM/DD format.
270 filter_end_date: Range end in YYYY/MM/DD format (default: today).
271 pep_type: Restrict to "pep" or "pop" (default: all).
273 Returns:
274 Number of matching projects.
275 """
276 if admin is None:
277 admin = []
278 statement = select(func.count()).select_from(Projects)
279 statement = self._add_condition(
280 statement,
281 namespace=namespace,
282 search_str=search_str,
283 admin_list=admin,
284 )
285 statement = self._add_date_filter_if_provided(
286 statement, filter_by, filter_start_date, filter_end_date
287 )
288 if pep_type:
289 statement = statement.where(Projects.pop.is_(pep_type == "pop"))
290 if tag:
291 statement = statement.where(Projects.tag == tag)
292 result = self._pep_db_engine.session_execute(statement).first()
294 try:
295 return result[0]
296 except IndexError:
297 return 0
299 def _get_projects(
300 self,
301 namespace: str = None,
302 tag: str = None,
303 search_str: str = None,
304 admin: str | list[str] | None = None,
305 limit: int = DEFAULT_LIMIT,
306 offset: int = DEFAULT_OFFSET,
307 order_by: str = "update_date",
308 order_desc: bool = False,
309 filter_by: Literal["submission_date", "last_update_date"] | None = None,
310 filter_start_date: str | None = None,
311 filter_end_date: str | None = None,
312 pep_type: Literal["pep", "pop"] | None = None,
313 ) -> list[AnnotationModel]:
314 """Get projects matching the given filters.
316 Args:
317 namespace: Namespace to restrict the search to.
318 tag: Exact tag to match.
319 search_str: String to search in name, tag, and description.
320 admin: Namespace(s) with admin/private access.
321 limit: Maximum number of results.
322 offset: Number of results to skip.
323 order_by: Sort field — "update_date", "name", "submission_date", or "stars".
324 order_desc: Sort in descending order if True.
325 filter_by: Date field for range filter — "submission_date" or "last_update_date".
326 filter_start_date: Range start in YYYY/MM/DD format.
327 filter_end_date: Range end in YYYY/MM/DD format (default: today).
328 pep_type: Restrict to "pep" or "pop" (default: all).
330 Returns:
331 List of AnnotationModel objects.
332 """
333 _LOGGER.info(
334 f"Running annotation search: (namespace: {namespace}, query: {search_str}."
335 )
337 if admin is None:
338 admin = []
339 statement = select(Projects)
341 statement = self._add_condition(
342 statement,
343 namespace=namespace,
344 search_str=search_str,
345 admin_list=admin,
346 tag=tag,
347 )
348 statement = self._add_date_filter_if_provided(
349 statement, filter_by, filter_start_date, filter_end_date
350 )
351 statement = statement
352 statement = self._add_order_by_keyword(statement, by=order_by, desc=order_desc)
353 statement = statement.limit(limit).offset(offset)
354 if pep_type:
355 statement = statement.where(Projects.pop.is_(pep_type == "pop"))
357 results_list = []
358 with Session(self._sa_engine) as session:
359 # Unique should be called because of the join with schema_mapping
360 results = session.scalars(statement).unique()
361 for result in results:
362 results_list.append(
363 AnnotationModel(
364 namespace=result.namespace,
365 name=result.name,
366 tag=result.tag,
367 is_private=result.private,
368 description=result.description,
369 number_of_samples=result.number_of_samples,
370 submission_date=str(result.submission_date),
371 last_update_date=str(result.last_update_date),
372 digest=result.digest,
373 pep_schema=(
374 f"{result.schema_mapping.schema_mapping.namespace}/{result.schema_mapping.schema_mapping.name}:{result.schema_mapping.version}"
375 if result.schema_mapping
376 else None
377 ),
378 pop=result.pop,
379 stars_number=result.number_of_stars,
380 forked_from=(
381 f"{result.forked_from_mapping.namespace}/{result.forked_from_mapping.name}:{result.forked_from_mapping.tag}"
382 if result.forked_from_id
383 else None
384 ),
385 )
386 )
387 return results_list
389 @staticmethod
390 def _add_order_by_keyword(
391 statement: Select, by: str = "update_date", desc: bool = False
392 ) -> Select:
393 """Add an ORDER BY clause to a SELECT statement.
395 Args:
396 statement: SQLAlchemy SELECT statement to augment.
397 by: Sort field — "name", "update_date", "submission_date", or "stars".
398 desc: Sort in descending order if True.
400 Returns:
401 Statement with ORDER BY applied.
402 """
403 if by == "update_date":
404 order_by_obj = Projects.last_update_date
405 elif by == "name":
406 order_by_obj = Projects.name
407 elif by == SUBMISSION_DATE_KEY:
408 order_by_obj = Projects.submission_date
409 elif by == "stars":
410 order_by_obj = Projects.number_of_stars
411 else:
412 _LOGGER.warning(
413 f"order by: '{by}' statement is unavailable. Projects are sorted by 'update_date'"
414 )
415 order_by_obj = Projects.last_update_date
417 if desc and by == "name":
418 order_by_obj = order_by_obj.desc()
420 elif by != "name" and not desc:
421 order_by_obj = order_by_obj.desc()
423 return statement.order_by(order_by_obj)
425 @staticmethod
426 def _add_condition(
427 statement: Select,
428 namespace: str = None,
429 search_str: str = None,
430 admin_list: str | list[str] | None = None,
431 tag: str = None,
432 ) -> Select:
433 """Add a WHERE clause to a project search statement.
435 Args:
436 statement: SQLAlchemy SELECT statement to augment.
437 namespace: Filter to this namespace.
438 search_str: String to search in name, tag, and description.
439 admin_list: Namespace(s) with admin/private access.
440 tag: Exact tag to match.
442 Returns:
443 Statement with WHERE clause applied.
444 """
445 admin_list = tuple_converter(admin_list)
446 if search_str:
447 sql_search_str = f"%{search_str}%"
448 search_query = or_(
449 Projects.name.ilike(sql_search_str),
450 Projects.tag.ilike(sql_search_str),
451 Projects.description.ilike(sql_search_str),
452 )
453 statement = statement.where(search_query)
454 if namespace:
455 statement = statement.where(Projects.namespace == namespace)
457 if tag:
458 statement = statement.where(Projects.tag == tag)
460 statement = statement.where(
461 or_(Projects.private.is_(False), Projects.namespace.in_(admin_list))
462 )
464 return statement
466 @staticmethod
467 def _add_date_filter_if_provided(
468 statement: Select,
469 filter_by: Literal["submission_date", "last_update_date"] | None,
470 filter_start_date: str | None,
471 filter_end_date: str | None = None,
472 ) -> Select:
473 """Add a date range filter to a SELECT statement.
475 Args:
476 statement: SQLAlchemy SELECT statement to augment.
477 filter_by: Date field to filter on — "submission_date" or "last_update_date".
478 filter_start_date: Range start in YYYY/MM/DD format.
479 filter_end_date: Range end in YYYY/MM/DD format (default: today).
481 Returns:
482 Statement with date filter applied.
483 """
484 if filter_by and filter_start_date:
485 start_date = convert_date_string_to_date(filter_start_date)
486 if filter_end_date:
487 end_date = convert_date_string_to_date(filter_end_date)
488 else:
489 end_date = datetime.now()
490 if filter_by == SUBMISSION_DATE_KEY:
491 statement = statement.filter(
492 Projects.submission_date.between(start_date, end_date)
493 )
494 elif filter_by == LAST_UPDATE_DATE_KEY:
495 statement = statement.filter(
496 Projects.last_update_date.between(start_date, end_date)
497 )
498 else:
499 raise FilterError("Invalid filter_by was provided!")
500 return statement
501 else:
502 if filter_by:
503 _LOGGER.warning(
504 "filter_start_date was not provided, skipping filter..."
505 )
506 return statement
508 def get_project_number_in_namespace(
509 self,
510 namespace: str,
511 admin: str | list[str] | None = None,
512 ) -> int:
513 """Get the number of projects in a namespace.
515 Args:
516 namespace: Namespace to count projects in.
517 admin: Namespace(s) with admin/private access.
519 Returns:
520 Number of accessible projects in the namespace.
521 """
522 if admin is None:
523 admin = []
524 statement = (
525 select(func.count())
526 .select_from(Projects)
527 .where(Projects.namespace == namespace)
528 )
529 statement = statement.where(
530 or_(Projects.private.is_(False), Projects.namespace.in_(admin))
531 )
533 result = self._pep_db_engine.session_execute(statement).first()
535 try:
536 return result[0]
537 except IndexError:
538 return 0
540 def get_by_rp_list(
541 self,
542 registry_paths: list[str],
543 admin: str | list[str] | None = None,
544 ) -> AnnotationList:
545 """Get project annotations for a list of registry paths in a single query.
547 Args:
548 registry_paths: List of registry paths ("namespace/name:tag").
549 admin: Namespace(s) where the caller has admin rights.
551 Returns:
552 AnnotationList preserving the input order (missing projects returned as None).
553 """
554 admin_tuple = tuple_converter(admin)
556 if isinstance(registry_paths, list):
557 or_statement_list = []
558 for path in registry_paths:
559 try:
560 namespace, name, tag = registry_path_converter(path)
561 or_statement_list.append(
562 and_(
563 Projects.name == name,
564 Projects.namespace == namespace,
565 Projects.tag == tag,
566 or_(
567 Projects.namespace.in_(admin_tuple),
568 Projects.private.is_(False),
569 ),
570 )
571 )
572 except RegistryPathError as err:
573 _LOGGER.error(str(err), registry_paths)
574 continue
575 if not or_statement_list:
576 _LOGGER.error("No valid registry paths were provided!")
577 return AnnotationList(
578 count=0,
579 limit=len(registry_paths),
580 offset=0,
581 results=[],
582 )
584 statement = select(Projects).where(or_(*or_statement_list))
585 anno_results = []
586 with Session(self._sa_engine) as session:
587 query_result = session.scalars(statement).unique()
588 for result in query_result:
589 project_obj = result
590 annot = AnnotationModel(
591 namespace=project_obj.namespace,
592 name=project_obj.name,
593 tag=project_obj.tag,
594 is_private=project_obj.private,
595 description=project_obj.description,
596 number_of_samples=project_obj.number_of_samples,
597 submission_date=str(project_obj.submission_date),
598 last_update_date=str(project_obj.last_update_date),
599 digest=project_obj.digest,
600 pep_schema=(
601 f"{project_obj.schema_mapping.schema_mapping.namespace}/{project_obj.schema_mapping.schema_mapping.name}:{project_obj.schema_mapping.version}"
602 if project_obj.schema_mapping
603 else None
604 ),
605 pop=project_obj.pop,
606 stars_number=project_obj.number_of_stars,
607 forked_from=(
608 f"{project_obj.forked_from_mapping.namespace}/{project_obj.forked_from_mapping.name}:{project_obj.forked_from_mapping.tag}"
609 if project_obj.forked_from_mapping
610 else None
611 ),
612 )
613 anno_results.append(annot)
615 found_dict = {f"{r.namespace}/{r.name}:{r.tag}": r for r in anno_results}
616 end_results = [found_dict.get(project) for project in registry_paths]
618 return_len = len(anno_results)
619 return AnnotationList(
620 count=return_len,
621 limit=len(registry_paths),
622 offset=0,
623 results=end_results,
624 )
626 else:
627 return self.get_by_rp(registry_paths, admin)
629 def get_projects_list(
630 self,
631 namespace: str = None,
632 search_str: str = None,
633 admin: str | list[str] | None = None,
634 limit: int = DEFAULT_LIMIT,
635 offset: int = DEFAULT_OFFSET,
636 order_by: str = "update_date",
637 order_desc: bool = False,
638 filter_by: Literal["submission_date", "last_update_date"] | None = None,
639 filter_start_date: str | None = None,
640 filter_end_date: str | None = None,
641 pep_type: Literal["pep", "pop"] | None = None,
642 ) -> list[RegistryPath]:
643 """Retrieve a list of registry paths matching the given filters.
645 Lightweight alternative to get() — returns only registry paths, no annotation data.
647 Args:
648 namespace: Namespace to restrict the search to.
649 search_str: String to search in name, tag, and description.
650 admin: Namespace(s) with admin/private access.
651 limit: Maximum number of results.
652 offset: Number of results to skip.
653 order_by: Sort field — "name", "update_date", "submission_date", or "stars".
654 order_desc: Sort in descending order if True.
655 filter_by: Date field for range filter — "submission_date" or "last_update_date".
656 filter_start_date: Range start in YYYY/MM/DD format.
657 filter_end_date: Range end in YYYY/MM/DD format (default: today).
658 pep_type: Restrict to "pep" or "pop" (default: all).
660 Returns:
661 List of RegistryPath objects.
662 """
663 _LOGGER.info(
664 f"Running project search: (namespace: {namespace}, query: {search_str}."
665 )
667 if admin is None:
668 admin = []
669 statement = select(Projects.namespace, Projects.name, Projects.tag)
671 statement = self._add_condition(
672 statement,
673 namespace=namespace,
674 search_str=search_str,
675 admin_list=admin,
676 )
677 statement = self._add_date_filter_if_provided(
678 statement, filter_by, filter_start_date, filter_end_date
679 )
680 statement = self._add_order_by_keyword(statement, by=order_by, desc=order_desc)
681 statement = statement.limit(limit).offset(offset)
682 if pep_type:
683 statement = statement.where(Projects.pop.is_(pep_type == "pop"))
685 results_list = []
686 with Session(self._sa_engine) as session:
687 results = session.execute(statement)
689 for result in results:
690 results_list.append(
691 RegistryPath(namespace=result[0], name=result[1], tag=result[2])
692 )
693 return results_list