PB-2094: Use CustomBaseCommand for all management commands

Author: msom

app/stac_api/management/commands/dummy_asset.py

@@ -1,19 +1,16 @@
 import hashlib
-import logging
 import random
 import uuid
 from io import BytesIO
 
 from django.conf import settings
-from django.core.management.base import BaseCommand
 
 from stac_api.utils import CommandHandler
+from stac_api.utils import CustomBaseCommand
 from stac_api.utils import get_s3_resource
 from stac_api.utils import get_sha256_multihash
 from stac_api.validators import MEDIA_TYPES
 
-logger = logging.getLogger(__name__)
-
 PREFIX = 'dummy-obj-'
 
 
@@ -66,7 +63,7 @@ def upload(self):
         self.print_success('Done')
 
 
-class Command(BaseCommand):
+class Command(CustomBaseCommand):
     help = f"""Upload dummy asset file on S3 for testing.
 
     The command upload dummy asset file with random data on S3 for testing.
@@ -80,6 +77,7 @@ class Command(BaseCommand):
     """
 
     def add_arguments(self, parser):
+        super().add_arguments(parser)
         parser.add_argument(
             'action',
             type=str,

app/stac_api/management/commands/dummy_asset_upload.py

@@ -1,20 +1,17 @@
-import logging
 import os
 
 from django.conf import settings
-from django.core.management.base import BaseCommand
 
 from stac_api.models.general import BaseAssetUpload
 from stac_api.models.item import Asset
 from stac_api.models.item import AssetUpload
 from stac_api.s3_multipart_upload import MultipartUpload
 from stac_api.utils import AVAILABLE_S3_BUCKETS
 from stac_api.utils import CommandHandler
+from stac_api.utils import CustomBaseCommand
 from stac_api.utils import get_asset_path
 from stac_api.utils import get_sha256_multihash
 
-logger = logging.getLogger(__name__)
-
 
 class DummyAssetUploadHandler(CommandHandler):
 
@@ -59,7 +56,7 @@ def start(self):
 
     def list(self):
         for upload in AssetUpload.objects.filter(status=BaseAssetUpload.Status.IN_PROGRESS):
-            print(f"> {upload.upload_id} (asset: {upload.asset.name})")
+            self.print(f"> {upload.upload_id} (asset: {upload.asset.name})")
 
     def complete(self):
         try:
@@ -78,11 +75,12 @@ def abort(self):
             self.print_error(f"upload_id {self.options['upload_id']} doesn't exist")
 
 
-class Command(BaseCommand):
+class Command(CustomBaseCommand):
     help = """Start dummy Multipart upload for asset file on S3 for testing.
     """
 
     def add_arguments(self, parser):
+        super().add_arguments(parser)
 
         subparsers = parser.add_subparsers(
             dest='action',

app/stac_api/management/commands/dummy_data.py

@@ -1,5 +1,4 @@
 import datetime
-import logging
 import random
 import string
 import time
@@ -11,16 +10,14 @@
 
 from django.contrib.gis.geos import Polygon
 from django.core.files.uploadedfile import SimpleUploadedFile
-from django.core.management.base import BaseCommand
 
 from stac_api.models.collection import Collection
 from stac_api.models.item import Asset
 from stac_api.models.item import Item
 from stac_api.utils import CommandHandler
+from stac_api.utils import CustomBaseCommand
 from stac_api.validators import MEDIA_TYPES
 
-logger = logging.getLogger(__name__)
-
 # Min/Max extent (roughly) of CH in LV95
 XMIN = 2570000
 XMAX = 2746000
@@ -252,7 +249,7 @@ def create_asset(self, item, asset_id):
         self.print('Asset %s/%s/%s created', item.collection.name, item.name, asset_id, level=3)
 
 
-class Command(BaseCommand):
+class Command(CustomBaseCommand):
     help = """Manage dummy data for performance testing.
 
     The command populates the database by default with
@@ -264,6 +261,7 @@ class Command(BaseCommand):
     """
 
     def add_arguments(self, parser):
+        super().add_arguments(parser)
         parser.add_argument(
             'action',
             type=str,

app/stac_api/management/commands/list_asset_uploads.py

@@ -1,17 +1,14 @@
 import json
-import logging
 
-from django.core.management.base import BaseCommand
 from django.core.serializers.json import DjangoJSONEncoder
 
 from stac_api.models.item import AssetUpload
 from stac_api.s3_multipart_upload import MultipartUpload
 from stac_api.serializers.upload import AssetUploadSerializer
 from stac_api.utils import CommandHandler
+from stac_api.utils import CustomBaseCommand
 from stac_api.utils import get_asset_path
 
-logger = logging.getLogger(__name__)
-
 
 class ListAssetUploadsHandler(CommandHandler):
 
@@ -98,7 +95,7 @@ def are_uploads_equal(s3_upload, db_upload):
         elif self.options['s3_only']:
             only_s3_uploads = s3_uploads
 
-        print(
+        self.print(
             json.dumps(
                 {
                     'uploads': uploads,
@@ -119,7 +116,7 @@ def are_uploads_equal(s3_upload, db_upload):
         )
 
 
-class Command(BaseCommand):
+class Command(CustomBaseCommand):
     help = """List all asset uploads object (DB and/or S3)
 
     This checks for all asset uploads object in DB (by default only returning the `in-progress`
@@ -138,6 +135,7 @@ class Command(BaseCommand):
 
     def add_arguments(self, parser):
         self.prog = parser.prog  # pylint: disable=attribute-defined-outside-init
+        super().add_arguments(parser)
 
         parser.add_argument(
             '--status',

app/stac_api/management/commands/populate_testdb.py

@@ -1,15 +1,13 @@
-import logging
 import os
 
 from django.conf import settings
-from django.core.management.base import BaseCommand
 
 from stac_api.sample_data import importer
 from stac_api.utils import CommandHandler
+from stac_api.utils import CustomBaseCommand
 
 # path definition relative to the directory that contains manage.py
 DATADIR = settings.BASE_DIR / 'app/stac_api/sample_data/'
-logger = logging.getLogger(__name__)
 
 
 class Handler(CommandHandler):
@@ -25,7 +23,7 @@ def populate(self):
         self.print_success('Done')
 
 
-class Command(BaseCommand):
+class Command(CustomBaseCommand):
     help = """Populates the local test database with sample data
 
     The sample data has to be located in stac_api/management/sample_data and

app/stac_api/management/commands/profile_cursor_paginator.py

@@ -1,19 +1,16 @@
 import cProfile
-import logging
 import os
 import pstats
 
 from django.conf import settings
-from django.core.management.base import BaseCommand
 
 from rest_framework.pagination import CursorPagination
 from rest_framework.request import Request
 from rest_framework.test import APIRequestFactory
 
 from stac_api.models.item import Item
 from stac_api.utils import CommandHandler
-
-logger = logging.getLogger(__name__)
+from stac_api.utils import CustomBaseCommand
 
 STAC_BASE_V = f'{settings.STAC_BASE}/v1'
 
@@ -44,7 +41,7 @@ def profiling(self):
         self.print_success('Done')
 
 
-class Command(BaseCommand):
+class Command(CustomBaseCommand):
     help = """Paginator paginate_queryset() profiling command
 
     Profiling of the method paginator.paginate_queryset(qs, request)
@@ -53,6 +50,7 @@ class Command(BaseCommand):
     """
 
     def add_arguments(self, parser):
+        super().add_arguments(parser)
         parser.add_argument(
             '--collection',
             type=str,

app/stac_api/management/commands/profile_item_serializer.py

@@ -1,17 +1,14 @@
 import cProfile
-import logging
 import os
 import pstats
 
 from django.conf import settings
-from django.core.management.base import BaseCommand
 
 from rest_framework.test import APIRequestFactory
 
 from stac_api.models.item import Item
 from stac_api.utils import CommandHandler
-
-logger = logging.getLogger(__name__)
+from stac_api.utils import CustomBaseCommand
 
 STAC_BASE_V = f'{settings.STAC_BASE}/v1'
 
@@ -40,7 +37,7 @@ def profiling(self):
         self.print_success('Done')
 
 
-class Command(BaseCommand):
+class Command(CustomBaseCommand):
     help = """ItemSerializer profiling command
 
     Profiling of the serialization of many items.
@@ -49,6 +46,7 @@ class Command(BaseCommand):
     """
 
     def add_arguments(self, parser):
+        super().add_arguments(parser)
         parser.add_argument(
             '--collection',
             type=str,

app/stac_api/management/commands/profile_serializer_vs_no_drf.py

@@ -1,16 +1,13 @@
 import json
-import logging
 from timeit import timeit
 
 from django.conf import settings
-from django.core.management.base import BaseCommand
 
 from rest_framework.test import APIRequestFactory
 
 from stac_api.models.item import Item
 from stac_api.utils import CommandHandler
-
-logger = logging.getLogger(__name__)
+from stac_api.utils import CustomBaseCommand
 
 STAC_BASE_V = f'{settings.STAC_BASE}/v1'
 
@@ -85,7 +82,7 @@ def serialize(qs):
         self.print_success('NO DRF time: %fms', no_drf_time / self.options['repeat'] * 1000)
 
 
-class Command(BaseCommand):
+class Command(CustomBaseCommand):
     help = """ItemSerializer vs simple serializer profiling command
 
     Profiling of the serialization of many items using DRF vs using a simple function.
@@ -94,6 +91,7 @@ class Command(BaseCommand):
     """
 
     def add_arguments(self, parser):
+        super().add_arguments(parser)
         parser.add_argument(
             '--collection',
             type=str,

app/stac_api/management/commands/reset_counter_tables.py

@@ -1,9 +1,9 @@
 import time
 
-from django.core.management.base import BaseCommand
 from django.db import connection
 
 from stac_api.utils import CommandHandler
+from stac_api.utils import CustomBaseCommand
 
 
 class Handler(CommandHandler):
@@ -59,7 +59,7 @@ def run(self):
         )
 
 
-class Command(BaseCommand):
+class Command(CustomBaseCommand):
     help = """Reset the summary counter tables.
 
     Truncates all the summary counter tables and repopulates with current data to make sure they are

app/stac_api/management/commands/update_asset_file_size.py

@@ -1,13 +1,11 @@
 import logging
 
-from django.core.management.base import BaseCommand
 from django.core.management.base import CommandParser
 
 from stac_api.models.collection import CollectionAsset
 from stac_api.models.item import Asset
 from stac_api.utils import CommandHandler
-
-logger = logging.getLogger(__name__)
+from stac_api.utils import CustomBaseCommand
 
 # increase the log level so boto3 doesn't spam the output
 logging.getLogger('boto3').setLevel(logging.WARNING)
@@ -17,15 +15,15 @@
 class Handler(CommandHandler):
 
     def update(self):
-        self.print_success('Running command to update file size')
+        self.print('Running command to update file size')
 
         asset_limit = self.options['count']
 
         asset_qs = Asset.objects.filter(file_size=0, is_external=False)
         total_asset_count = asset_qs.count()
         assets = asset_qs.all()[:asset_limit]
 
-        self.print_success(f'Update file size for {len(assets)} assets out of {total_asset_count}')
+        self.print(f'Update file size for {len(assets)} assets out of {total_asset_count}')
 
         for asset in assets:
             try:
@@ -45,14 +43,13 @@ def update(self):
                 asset.file_size = None
                 asset.save()
                 print("_", end="", flush=True)
-                logger.error('file %s could not be found', asset.file)
-        print()
+                self.print_error('file %s could not be found', asset.file)
 
         collection_asset_qs = CollectionAsset.objects.filter(file_size=0)
         total_asset_count = collection_asset_qs.count()
         collection_assets = collection_asset_qs.all()[:asset_limit]
 
-        self.print_success(
+        self.print(
             f"Update file size for {len(collection_assets)} collection assets out of "
             f"{total_asset_count}"
         )
@@ -73,14 +70,12 @@ def update(self):
                 # bucket.
                 collection_asset.file_size = None
                 collection_asset.save()
-                print("_", end="", flush=True)
-                logger.error('file %s could not be found', collection_asset.file)
+                self.print_error('file %s could not be found', collection_asset.file)
 
-        print()
         self.print_success('Update completed')
 
 
-class Command(BaseCommand):
+class Command(CustomBaseCommand):
     help = """Requests the file size of every asset / collection asset from the s3 bucket and
         updates the value in the database"""