Skip to content
Snippets Groups Projects
import_files.py 7.99 KiB
Newer Older
  • Learn to ignore specific revisions
  • from django.conf import settings
    
    from django.core.files import File
    
    from django.core.management.base import BaseCommand, CommandError
    
    Eliot Berriot's avatar
    Eliot Berriot committed
    from funkwhale_api.music import models, tasks
    
    from funkwhale_api.users.models import User
    
    Eliot Berriot's avatar
    Eliot Berriot committed
        help = "Import audio files mathinc given glob pattern"
    
            parser.add_argument("path", nargs="+", type=str)
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                "--recursive",
                action="store_true",
                dest="recursive",
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                help="Will match the pattern recursively (including subdirectories)",
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                "--username",
                dest="username",
                help="The username of the user you want to be bound to the import",
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                "--async",
                action="store_true",
                dest="async",
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                help="Will launch celery tasks for each file to import instead of doing it synchronously and block the CLI",
    
            parser.add_argument(
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                "--exit",
                "-x",
                action="store_true",
                dest="exit_on_failure",
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                help="Use this flag to disable error catching",
    
            )
            parser.add_argument(
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                "--in-place",
                "-i",
                action="store_true",
                dest="in_place",
    
                default=False,
                help=(
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    "Import files without duplicating them into the media directory."
                    "For in-place import to work, the music files must be readable"
                    "by the web-server and funkwhale api and celeryworker processes."
                    "You may want to use this if you have a big music library to "
                    "import and not much disk space available."
                ),
    
            parser.add_argument(
                "--update",
                action="store_true",
                dest="update",
                default=False,
                help=(
                    "Use this flag to replace duplicates (tracks with same "
                    "musicbrainz mbid) on import with their newest version."
                ),
            )
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                "--noinput",
                "--no-input",
                action="store_false",
                dest="interactive",
    
                help="Do NOT prompt the user for input of any kind.",
            )
    
        def handle(self, *args, **options):
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            if options["recursive"]:
                glob_kwargs["recursive"] = True
    
                for import_path in options["path"]:
                    matching += glob.glob(import_path, **glob_kwargs)
                matching = sorted(list(set(matching)))
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                raise Exception("You need Python 3.5 to use the --recursive flag")
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            if options["in_place"]:
    
                self.stdout.write(
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    "Checking imported paths against settings.MUSIC_DIRECTORY_PATH"
                )
    
                p = settings.MUSIC_DIRECTORY_PATH
                if not p:
                    raise CommandError(
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                        "Importing in-place requires setting the "
                        "MUSIC_DIRECTORY_PATH variable"
                    )
    
                for m in matching:
                    if not m.startswith(p):
                        raise CommandError(
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                            "Importing in-place only works if importing"
                            "from {} (MUSIC_DIRECTORY_PATH), as this directory"
                            "needs to be accessible by the webserver."
                            "Culprit: {}".format(p, m)
                        )
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                raise CommandError("No file matching pattern, aborting")
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            if options["username"]:
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    user = User.objects.get(username=options["username"])
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    raise CommandError("Invalid username")
    
            else:
                # we bind the import to the first registered superuser
                try:
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    user = User.objects.filter(is_superuser=True).order_by("pk").first()
    
                    assert user is not None
                except AssertionError:
                    raise CommandError(
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                        "No superuser available, please provide a --username"
                    )
    
            filtered = self.filter_matching(matching)
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            self.stdout.write("Import summary:")
            self.stdout.write(
                "- {} files found matching this pattern: {}".format(
                    len(matching), options["path"]
                )
            )
            self.stdout.write(
                "- {} files already found in database".format(len(filtered["skipped"]))
            )
            self.stdout.write("- {} new files".format(len(filtered["new"])))
    
            self.stdout.write(
                "Selected options: {}".format(
                    ", ".join(["in place" if options["in_place"] else "copy music files"])
                )
            )
            if len(filtered["new"]) == 0:
                self.stdout.write("Nothing new to import, exiting")
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            if options["interactive"]:
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    "Are you sure you want to do this?\n\n"
    
                    "Type 'yes' to continue, or 'no' to cancel: "
                )
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                if input("".join(message)) != "yes":
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            batch, errors = self.do_import(filtered["new"], user=user, options=options)
            message = "Successfully imported {} tracks"
            if options["async"]:
                message = "Successfully launched import for {} tracks"
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            self.stdout.write(message.format(len(filtered["new"])))
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                self.stderr.write("{} tracks could not be imported:".format(len(errors)))
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    self.stderr.write("- {}: {}".format(path, error))
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                "For details, please refer to import batch #{}".format(batch.pk)
            )
    
        def filter_matching(self, matching):
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            sources = ["file://{}".format(p) for p in matching]
    
            # we skip reimport for path that are already found
            # as a TrackFile.source
            existing = models.TrackFile.objects.filter(source__in=sources)
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            existing = existing.values_list("source", flat=True)
            existing = set([p.replace("file://", "", 1) for p in existing])
    
            skipped = set(matching) & existing
            result = {
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                "initial": matching,
                "skipped": list(sorted(skipped)),
                "new": list(sorted(set(matching) - skipped)),
    
            }
            return result
    
        def do_import(self, paths, user, options):
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            message = "{i}/{total} Importing {path}..."
            if options["async"]:
                message = "{i}/{total} Launching import for {path}..."
    
            # we create an import batch binded to the user
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            async = options["async"]
    
            import_handler = tasks.import_job_run.delay if async else tasks.import_job_run
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            batch = user.imports.create(source="shell")
    
            for i, path in list(enumerate(paths)):
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    self.stdout.write(message.format(path=path, i=i + 1, total=len(paths)))
    
                    self.import_file(path, batch, import_handler, options)
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    if options["exit_on_failure"]:
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    m = "Error while importing {}: {} {}".format(
                        path, e.__class__.__name__, e
                    )
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                    errors.append((path, "{} {}".format(e.__class__.__name__, e)))
    
    
        def import_file(self, path, batch, import_handler, options):
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            job = batch.jobs.create(source="file://" + path)
            if not options["in_place"]:
    
                name = os.path.basename(path)
    
    Eliot Berriot's avatar
    Eliot Berriot committed
                with open(path, "rb") as f:
    
                    job.audio_file.save(name, File(f))
    
    Eliot Berriot's avatar
    Eliot Berriot committed
            import_handler(import_job_id=job.pk, use_acoustid=False)