Fixing dangling databases (wtf) 🧖🏽‍♂️

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 1.12.1 : 2024-03-24
+
+- **Fixed**: Prune also inactive databases which are standing still for some reason
+
 ## 1.12.0 : 2024-03-05
 
 - **Added**: `Recreate queue` in the administration

apps/core/jobs.py

@@ -11,7 +11,6 @@
 from uuid import UUID
 
 import docker
-import psycopg
 import sentry_sdk
 from django.conf import settings
 from django.db import connection
@@ -52,6 +51,10 @@ def prepare(self):
             cursor.execute(f"CREATE DATABASE {self._database_name} OWNER {self._database_name};")
             connection.commit()
 
+        self._task.additional_information['database'] = {
+            'name': self._database_name
+        }
+
         # Recover database
         command = [settings.PSQL_PATH, self._database_name]
 
@@ -197,6 +200,7 @@ def cleanup(self):
         with connection.cursor() as cursor:
             cursor.execute(f"DROP DATABASE {self._database_name};")
             cursor.execute(f"DROP USER {self._database_name};")
+            connection.commit()
 
     @staticmethod
     def execute(task_id: UUID, public_only: bool) -> Optional[Task]:

apps/core/management/commands/prune.py

@@ -1,10 +1,13 @@
 import docker
 from django.core.management import BaseCommand
+from django.db import connection
 from django.utils import timezone
 
+from apps.core.models import Task
+
 
 class Command(BaseCommand):
-    help = "Prune unused docker containers/images"
+    help = "Prune unused docker containers/images and databases"
 
     def handle(self, *args, **options):
         started_at = timezone.now()
@@ -14,5 +17,11 @@ def handle(self, *args, **options):
         client.containers.prune()
         client.images.prune(filters={"dangling": False, "until": "5m"})
 
+        for task in Task.objects.filter(status=Task.Status.FAILED, additional_information__database__isnull=False):
+            with connection.cursor() as cursor:
+                cursor.execute(f"DROP DATABASE IF EXISTS {task.additional_information['database']['name']};")
+                cursor.execute(f"DROP USER IF EXISTS {task.additional_information['database']['name']};")
+                connection.commit()
+
         self.stdout.write(f"Finished: {timezone.now().isoformat()}")
         self.stdout.write(f"Duration: {timezone.now() - started_at}")

apps/core/migrations/0015_alter_task_additional_information.py

@@ -0,0 +1,18 @@
+# Generated by Django 5.0.3 on 2024-03-24 09:32
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("core", "0014_alter_task_image"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="task",
+            name="additional_information",
+            field=models.JSONField(default=dict),
+        ),
+    ]

apps/core/models/evaluation.py

@@ -47,8 +47,8 @@ def execute_tasks(sender, instance: Evaluation, created: bool, **kwargs):
                     assigment=instance.assignment,
                     executor=Task.Executor.EVALUATION,
                     image=row["Link"],
-                    additional_information=row,
                 )
+                task.additional_information['evaluation'] = row
                 instance.tasks.add(task)
                 jobs.append(
                     Queue.prepare_data(

apps/core/models/task.py

@@ -38,7 +38,7 @@ class Executor(models.TextChoices):
     message = models.TextField(null=True, editable=False)
     output = models.TextField(null=True, editable=False)
     note = models.TextField(null=True)
-    additional_information = models.JSONField(null=True)
+    additional_information = models.JSONField(default=dict)
 
     def get_absolute_url(self):
         return reverse("task-detail", kwargs={"task_id": self.pk})

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "dbs_tester"
-version = "1.12.0"
+version = "1.12.1"
 description = "DBS tester application"
 authors = [
     "Jakub Dubec <[email protected]>"