Merge pull request #1 from dex3r/main

Added auto_worker_type endpoint

.gitignore

@@ -113,6 +113,9 @@ venv/
 ENV/
 env.bak/
 venv.bak/
+.env_docker
+env_docker/
+env_docker.bak/
 
 # Spyder project settings
 .spyderproject

docker-compose.yaml

@@ -21,7 +21,9 @@ services:
       - redis
 
   postgres:
-    image: postgres:15.3-alpine
+    build:
+      context: docker/postgres
+      dockerfile: dockerfile
     container_name: postgres
     restart: always
     environment:

docker/postgres/dockerfile

@@ -0,0 +1,7 @@
+FROM postgres:15.3
+RUN apt-get update && apt-get install -y postgresql-15-cron
+
+RUN echo "shared_preload_libraries='pg_cron'" >> /usr/share/postgresql/postgresql.conf.sample
+RUN echo "cron.database_name='postgres'" >> /usr/share/postgresql/postgresql.conf.sample
+
+COPY init-db /docker-entrypoint-initdb.d

docker/postgres/init-db/002-pg-cron.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+
+# use same db as the one from env
+dbname="postgres"
+
+# create custom config
+customconf=/var/lib/postgresql/data/custom-conf.conf
+echo "" > $customconf
+echo "shared_preload_libraries = 'pg_cron'" >> $customconf
+echo "cron.database_name = '$dbname'" >> $customconf
+chown postgres $customconf
+chgrp postgres $customconf
+
+# include custom config from main config
+conf=/var/lib/postgresql/data/postgresql.conf
+found=$(grep "include = '$customconf'" $conf)
+if [ -z "$found" ]; then
+  echo "include = '$customconf'" >> $conf
+fi

docker/postgres/init-db/003-main.sql

@@ -0,0 +1 @@
+CREATE EXTENSION pg_cron;

horde/apis/models/v2.py

@@ -1549,3 +1549,15 @@ def __init__(self, api):
                 ),
             },
         )
+
+        self.response_model_auto_worker_type = api.model(
+            "AutoWorkerType",
+            {
+                "recommended_worker_type": fields.String(
+                    example="image",
+                    description="The recommended type of worker.",
+                    enum=["image", "text"],
+                    required=True,
+                ),
+            },
+        )

horde/apis/v2/__init__.py

@@ -52,3 +52,4 @@
 api.add_resource(base.DocsTerms, "/documents/terms")
 api.add_resource(base.DocsPrivacy, "/documents/privacy")
 api.add_resource(base.DocsSponsors, "/documents/sponsors")
+api.add_resource(base.AutoWorkerType, "/auto_worker_type")

horde/apis/v2/base.py

@@ -3070,3 +3070,56 @@ def get(self):
         if self.args.format == "markdown":
             return {"markdown": markdownify(html_template).strip("\n")}, 200
         return {"html": html_template}, 200
+
+class AutoWorkerType(Resource):
+    get_parser = reqparse.RequestParser()
+    get_parser.add_argument(
+        "apikey",
+        type=str,
+        required=True,
+        help="A User API key.",
+        location="headers",
+    )
+    get_parser.add_argument(
+        "Client-Agent",
+        default="unknown:0:unknown",
+        type=str,
+        required=False,
+        help="The client name and version.",
+        location="headers",
+    )
+
+    @api.expect(get_parser)
+    @api.marshal_with(
+        models.response_model_auto_worker_type,
+        code=200,
+        description="Recommended worker type details",
+        skip_none=True,
+    )
+    @api.response(400, "Validation Error", models.response_model_error)
+    @api.response(401, "Invalid API Key", models.response_model_error)
+    def get(self):
+        """Get the recommended worker type for the user"""
+        self.args = self.get_parser.parse_args()
+        self.user = database.find_user_by_api_key(self.args["apikey"])
+        if not self.user:
+            raise e.InvalidAPIKey("get auto worker type")
+
+        active_workers = database.get_active_workers()
+        image_workers_count = 0
+        text_workers_count = 0
+
+        for worker in active_workers:
+            if worker.worker_type == "image":
+                image_workers_count += 1
+            elif worker.worker_type == "text":
+                text_workers_count += 1
+            else:
+                logger.warning(f"Unknown worker type {worker.worker_type} for worker {worker.id}")
+
+        logger.debug(f"Image workers: {image_workers_count} Text workers: {text_workers_count}")
+
+        if image_workers_count > text_workers_count:
+            return {"recommended_worker_type": "text"}, 200
+        else:
+            return {"recommended_worker_type": "image"}, 200