Skip to content

Commit

Permalink
feat: allow seeking a worker by name
Browse files Browse the repository at this point in the history
  • Loading branch information
db0 committed Sep 15, 2024
1 parent c8a455b commit be9cb47
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 3 deletions.
17 changes: 14 additions & 3 deletions horde/apis/v2/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -768,6 +768,15 @@ class Workers(Resource):
location="args",
)

get_parser.add_argument(
"name",
required=False,
default=None,
type=str,
help="Find a worker by name (case insensitive).",
location="args",
)

@api.expect(get_parser)
@logger.catch(reraise=True)
# @cache.cached(timeout=10, query_string=True)
Expand Down Expand Up @@ -813,9 +822,11 @@ def get_worker_info_list(self, details_privilege):
return workers_ret

def parse_worker_by_query(self, workers_list):
if not self.args.type:
return workers_list
return [w for w in workers_list if w["type"] == self.args.type]
if self.args.name:
return [w for w in workers_list if w["name"].lower() == self.args.name.lower()]
if self.args.type:
return [w for w in workers_list if w["type"] == self.args.type]
return workers_list


class WorkerSingle(Resource):
Expand Down
7 changes: 7 additions & 0 deletions horde/database/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,13 @@ def find_worker_by_name(worker_name, worker_class=ImageWorker):
return worker


def find_worker_id_by_name(worker_name):
for worker_class in [ImageWorker, TextWorker, InterrogationWorker]:
worker_id = db.session.query(worker_class.id).filter_by(name=worker_name).first()
if worker_id:
return worker_id


def worker_name_exists(worker_name):
for worker_class in [ImageWorker, TextWorker, InterrogationWorker]:
worker = db.session.query(worker_class).filter_by(name=worker_name).count()
Expand Down

0 comments on commit be9cb47

Please sign in to comment.