def get_active_users( cls, company, project_ids: Sequence[str], user_ids: Optional[Sequence[str]] = None, ) -> Set[str]: """ Get the set of user ids that created tasks/models in the given projects If project_ids is empty then all projects are examined If user_ids are passed then only subset of these users is returned """ with TimingContext("mongo", "active_users_in_projects"): query = Q(company=company) if user_ids: query &= Q(user__in=user_ids) projects_query = query if project_ids: project_ids = _ids_with_children(project_ids) query &= Q(project__in=project_ids) projects_query &= Q(id__in=project_ids) res = set(Project.objects(projects_query).distinct(field="user")) for cls_ in (Task, Model): res |= set(cls_.objects(query).distinct(field="user")) return res
def _ids_with_parents(project_ids: Sequence[str]) -> Sequence[str]: """ Return project ids with all the parent projects """ projects = Project.objects(id__in=project_ids).only("id", "path") parent_ids = set( itertools.chain.from_iterable(p.path for p in projects if p.path)) return list({*(p.id for p in projects), *parent_ids})
def update_project_time(project_ids: Union[str, Sequence[str]]): if not project_ids: return if isinstance(project_ids, str): project_ids = [project_ids] return Project.objects(id__in=project_ids).update(last_update=datetime.utcnow())
def find_or_create( cls, user: str, company: str, project_name: str, description: str, project_id: str = None, tags: Sequence[str] = None, system_tags: Sequence[str] = None, default_output_destination: str = None, parent_creation_params: dict = None, ) -> str: """ Find a project named `project_name` or create a new one. Returns project ID """ if not project_id and not project_name: raise errors.bad_request.ValidationError( "project id or name required") if project_id: project = Project.objects(company=company, id=project_id).only("id").first() if not project: raise errors.bad_request.InvalidProjectId(id=project_id) return project_id project_name, _ = _validate_project_name(project_name) project = Project.objects(company=company, name=project_name).only("id").first() if project: return project.id return cls.create( user=user, company=company, name=project_name, description=description, tags=tags, system_tags=system_tags, default_output_destination=default_output_destination, parent_creation_params=parent_creation_params, )
def get_project_tags( cls, company_id: str, include_system: bool, projects: Sequence[str] = None, filter_: Dict[str, Sequence[str]] = None, ) -> Tuple[Sequence[str], Sequence[str]]: with TimingContext("mongo", "get_tags_from_db"): query = Q(company=company_id) if filter_: for name, vals in filter_.items(): if vals: query &= GetMixin.get_list_field_query(name, vals) if projects: query &= Q(id__in=_ids_with_children(projects)) tags = Project.objects(query).distinct("tags") system_tags = (Project.objects(query).distinct("system_tags") if include_system else []) return tags, system_tags
def _get_writable_project_from_name( company, name, _only: Optional[Sequence[str]] = ("id", "name", "path", "company", "parent"), ) -> Optional[Project]: """ Return a project from name. If the project not found then return None """ qs = Project.objects(company=company, name=name) if _only: qs = qs.only(*_only) return qs.first()
def merge_project(cls, company, source_id: str, destination_id: str) -> Tuple[int, int, Set[str]]: """ Move all the tasks and sub projects from the source project to the destination Remove the source project Return the amounts of moved entities and subprojects + set of all the affected project ids """ with TimingContext("mongo", "move_project"): if source_id == destination_id: raise errors.bad_request.ProjectSourceAndDestinationAreTheSame( source=source_id) source = Project.get(company, source_id) destination = Project.get(company, destination_id) if source_id in destination.path: raise errors.bad_request.ProjectCannotBeMergedIntoItsChild( source=source_id, destination=destination_id) children = _get_sub_projects([source.id], _only=("id", "name", "parent", "path"))[source.id] cls.validate_projects_depth( projects=children, old_parent_depth=len(source.path) + 1, new_parent_depth=len(destination.path) + 1, ) moved_entities = 0 for entity_type in (Task, Model): moved_entities += entity_type.objects( company=company, project=source_id, system_tags__nin=[EntityVisibility.archived.value], ).update(upsert=False, project=destination_id) moved_sub_projects = 0 for child in Project.objects(company=company, parent=source_id): _reposition_project_with_children( project=child, children=[c for c in children if c.parent == child.id], parent=destination, ) moved_sub_projects += 1 affected = {source.id, *(source.path or [])} source.delete() if destination: destination.update(last_update=datetime.utcnow()) affected.update({destination.id, *(destination.path or [])}) return moved_entities, moved_sub_projects, affected
def _update_task_name(task: Task): if not task or not task.project: return project = Project.objects(id=task.project).only("name").first() if not project: return _, _, name_prefix = project.name.rpartition("/") name_mask = re.compile(rf"{re.escape(name_prefix)}( #\d+)?$") count = Task.objects(project=task.project, system_tags__in=["pipeline"], name=name_mask).count() new_name = f"{name_prefix} #{count}" if count > 0 else name_prefix task.update(name=new_name)
def delete_project( company: str, project_id: str, force: bool, delete_contents: bool ) -> Tuple[DeleteProjectResult, Set[str]]: project = Project.get_for_writing( company=company, id=project_id, _only=("id", "path") ) if not project: raise errors.bad_request.InvalidProjectId(id=project_id) project_ids = _ids_with_children([project_id]) if not force: for cls, error in ( (Task, errors.bad_request.ProjectHasTasks), (Model, errors.bad_request.ProjectHasModels), ): non_archived = cls.objects( project__in=project_ids, system_tags__nin=[EntityVisibility.archived.value], ).only("id") if non_archived: raise error("use force=true to delete", id=project_id) if not delete_contents: with TimingContext("mongo", "update_children"): for cls in (Model, Task): updated_count = cls.objects(project__in=project_ids).update( project=None ) res = DeleteProjectResult(disassociated_tasks=updated_count) else: deleted_models, model_urls = _delete_models(projects=project_ids) deleted_tasks, event_urls, artifact_urls = _delete_tasks( company=company, projects=project_ids ) res = DeleteProjectResult( deleted_tasks=deleted_tasks, deleted_models=deleted_models, urls=TaskUrls( model_urls=list(model_urls), event_urls=list(event_urls), artifact_urls=list(artifact_urls), ), ) affected = {*project_ids, *(project.path or [])} res.deleted = Project.objects(id__in=project_ids).delete() return res, affected
def get_by_id(call): assert isinstance(call, APICall) project_id = call.data["project"] with translate_errors_context(): with TimingContext("mongo", "projects_by_id"): query = Q(id=project_id) & get_company_or_none_constraint( call.identity.company) project = Project.objects(query).first() if not project: raise errors.bad_request.InvalidProjectId(id=project_id) project_dict = project.to_proper_dict() conform_output_tags(call, project_dict) call.result.data = {"project": project_dict}
def _get_sub_projects( project_ids: Sequence[str], _only: Sequence[str] = ("id", "path") ) -> Mapping[str, Sequence[Project]]: """ Return the list of child projects of all the levels for the parent project ids """ qs = Project.objects(path__in=project_ids) if _only: _only = set(_only) | {"path"} qs = qs.only(*_only) subprojects = list(qs) return { pid: [s for s in subprojects if pid in (s.path or [])] for pid in project_ids }
def get_projects_with_active_user( cls, company: str, users: Sequence[str], project_ids: Optional[Sequence[str]] = None, allow_public: bool = True, ) -> Sequence[str]: """ Get the projects ids where user created any tasks including all the parents of these projects If project ids are specified then filter the results by these project ids """ query = Q(user__in=users) if allow_public: query &= get_company_or_none_constraint(company) else: query &= Q(company=company) user_projects_query = query if project_ids: ids_with_children = _ids_with_children(project_ids) query &= Q(project__in=ids_with_children) user_projects_query &= Q(id__in=ids_with_children) res = {p.id for p in Project.objects(user_projects_query).only("id")} for cls_ in (Task, Model): res |= set(cls_.objects(query).distinct(field="project")) res = list(res) if not res: return res ids_with_parents = _ids_with_parents(res) if project_ids: return [pid for pid in ids_with_parents if pid in project_ids] return ids_with_parents
def status_report( self, company_id: str, user_id: str, ip: str, report: StatusReportRequest, tags: Sequence[str] = None, ) -> None: """ Write worker status report :param company_id: worker's company ID :param user_id: user_id ID under which this worker is running :param ip: worker IP :param report: the report itself :param tags: tags for this worker :raise bad_request.InvalidTaskId: the reported task was not found :return: worker entry instance """ entry = self._get_worker(company_id, user_id, report.worker) try: entry.ip = ip now = datetime.utcnow() entry.last_activity_time = now if tags is not None: entry.tags = tags if report.machine_stats: self._log_stats_to_es( company_id=company_id, company_name=entry.company.name, worker=report.worker, timestamp=report.timestamp, task=report.task, machine_stats=report.machine_stats, ) entry.queue = report.queue if report.queues: entry.queues = report.queues if not report.task: entry.task = None entry.project = None else: with translate_errors_context(): query = dict(id=report.task, company=company_id) update = dict( last_worker=report.worker, last_worker_report=now, last_update=now, last_change=now, ) # modify(new=True, ...) returns the modified object task = Task.objects(**query).modify(new=True, **update) if not task: raise bad_request.InvalidTaskId(**query) entry.task = IdNameEntry(id=task.id, name=task.name) entry.project = None if task.project: project = Project.objects( id=task.project).only("name").first() if project: entry.project = IdNameEntry(id=project.id, name=project.name) entry.last_report_time = now except APIError: raise except Exception as e: msg = "Failed processing worker status report" log.exception(msg) raise server_error.DataError(msg, err=e.args[0]) finally: self._save_worker(entry)
def _ids_with_children(project_ids: Sequence[str]) -> Sequence[str]: """ Return project ids with the ids of all the subprojects """ subprojects = Project.objects(path__in=project_ids).only("id") return list({*project_ids, *(child.id for child in subprojects)})