Files @ d2d8513e900b
Branch filter:

Location: rattail-project/rattail/rattail/commands/importing.py

Lance Edgar
Fix dependency bug when testing coverage with tox
# -*- coding: utf-8 -*-
################################################################################
#
#  Rattail -- Retail Software Framework
#  Copyright © 2010-2016 Lance Edgar
#
#  This file is part of Rattail.
#
#  Rattail is free software: you can redistribute it and/or modify it under the
#  terms of the GNU Affero General Public License as published by the Free
#  Software Foundation, either version 3 of the License, or (at your option)
#  any later version.
#
#  Rattail is distributed in the hope that it will be useful, but WITHOUT ANY
#  WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
#  FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public License for
#  more details.
#
#  You should have received a copy of the GNU Affero General Public License
#  along with Rattail.  If not, see <http://www.gnu.org/licenses/>.
#
################################################################################
"""
Importing Commands
"""

from __future__ import unicode_literals, absolute_import

import logging

from rattail.commands.core import Subcommand, date_argument
from rattail.util import load_object


log = logging.getLogger(__name__)


class ImportSubcommand(Subcommand):
    """
    Base class for subcommands which use the (new) data importing system.
    """

    # TODO: move this into Subcommand or something..
    parent_name = None
    def __init__(self, *args, **kwargs):
        super(ImportSubcommand, self).__init__(*args, **kwargs)
        if self.parent:
            self.parent_name = self.parent.name

    def get_handler_factory(self):
        """
        Subclasses must override this, and return a callable that creates an
        import handler instance which the command should use.
        """
        raise NotImplementedError

    def get_handler(self, **kwargs):
        """
        Returns a handler instance to be used by the command.
        """
        factory = self.get_handler_factory()
        kwargs.setdefault('config', getattr(self, 'config', None))
        kwargs.setdefault('command', self)
        kwargs.setdefault('progress', self.progress)
        if 'args' in kwargs:
            args = kwargs['args']
            kwargs.setdefault('dry_run', args.dry_run)
            # kwargs.setdefault('max_create', args.max_create)
            # kwargs.setdefault('max_update', args.max_update)
            # kwargs.setdefault('max_delete', args.max_delete)
            # kwargs.setdefault('max_total', args.max_total)
        kwargs = self.get_handler_kwargs(**kwargs)
        return factory(**kwargs)

    def get_handler_kwargs(self, **kwargs):
        """
        Return a dict of kwargs to be passed to the handler factory.
        """
        return kwargs

    def add_parser_args(self, parser):

        # model names (aka importer keys)
        doc = ("Which data models to import.  If you specify any, then only "
               "data for those models will be imported.  If you do not specify "
               "any, then all *default* models will be imported.")
        try:
            handler = self.get_handler()
        except NotImplementedError:
            pass
        else:
            doc += "  Supported models are: ({})".format(', '.join(handler.get_importer_keys()))
        parser.add_argument('models', nargs='*', metavar='MODEL', help=doc)

        # start/end date
        parser.add_argument('--start-date', type=date_argument,
                            help="Optional (inclusive) starting point for date range, by which host "
                            "data should be filtered.  Only used by certain importers.")
        parser.add_argument('--end-date', type=date_argument,
                            help="Optional (inclusive) ending point for date range, by which host "
                            "data should be filtered.  Only used by certain importers.")

        # allow create?
        parser.add_argument('--create', action='store_true', default=True,
                            help="Allow new records to be created during the import.")
        parser.add_argument('--no-create', action='store_false', dest='create',
                            help="Do not allow new records to be created during the import.")
        parser.add_argument('--max-create', type=int, metavar='COUNT',
                            help="Maximum number of records which may be created, after which a "
                            "given import task should stop.  Note that this applies on a per-model "
                            "basis and not overall.")

        # allow update?
        parser.add_argument('--update', action='store_true', default=True,
                            help="Allow existing records to be updated during the import.")
        parser.add_argument('--no-update', action='store_false', dest='update',
                            help="Do not allow existing records to be updated during the import.")
        parser.add_argument('--max-update', type=int, metavar='COUNT',
                            help="Maximum number of records which may be updated, after which a "
                            "given import task should stop.  Note that this applies on a per-model "
                            "basis and not overall.")

        # allow delete?
        parser.add_argument('--delete', action='store_true', default=False,
                            help="Allow records to be deleted during the import.")
        parser.add_argument('--no-delete', action='store_false', dest='delete',
                            help="Do not allow records to be deleted during the import.")
        parser.add_argument('--max-delete', type=int, metavar='COUNT',
                            help="Maximum number of records which may be deleted, after which a "
                            "given import task should stop.  Note that this applies on a per-model "
                            "basis and not overall.")

        # max total changes, per model
        parser.add_argument('--max-total', type=int, metavar='COUNT',
                            help="Maximum number of *any* record changes which may occur, after which "
                            "a given import task should stop.  Note that this applies on a per-model "
                            "basis and not overall.")

        # treat changes as warnings?
        parser.add_argument('--warnings', '-W', action='store_true',
                            help="Set this flag if you expect a \"clean\" import, and wish for any "
                            "changes which do occur to be processed further and/or specially.  The "
                            "behavior of this flag is ultimately up to the import handler, but the "
                            "default is to send an email notification.")

        # dry run?
        parser.add_argument('--dry-run', action='store_true',
                            help="Go through the full motions and allow logging etc. to "
                            "occur, but rollback (abort) the transaction at the end.")

    def run(self, args):
        log.info("begin `{} {}` for data models: {}".format(
                self.parent_name, self.name, ', '.join(args.models or ["(ALL)"])))

        handler = self.get_handler(args=args)
        models = args.models or handler.get_default_keys()
        log.debug("using handler: {}".format(handler))
        log.debug("importing models: {}".format(models))
        log.debug("args are: {}".format(args))

        kwargs = {
            'dry_run': args.dry_run,
            'warnings': args.warnings,
            'max_create': args.max_create,
            'max_update': args.max_update,
            'max_delete': args.max_delete,
            'max_total': args.max_total,
            'progress': self.progress,
        }
        handler.import_data(*models, **kwargs)

        # TODO: should this logging happen elsewhere / be customizable?
        if args.dry_run:
            log.info("dry run, so transaction was rolled back")
        else:
            log.info("transaction was committed")


class ImportRattail(ImportSubcommand):
    """
    Import data from another Rattail database
    """
    name = 'import-rattail'
    description = __doc__.strip()
    handler_key = 'rattail'
    default_handler_spec = 'rattail.importing.rattail:FromRattailToRattail'

    def get_handler_factory(self):
        spec = self.config.get('rattail.importing', '{}.handler'.format(self.handler_key),
                               default=self.default_handler_spec)
        return load_object(spec)

    def add_parser_args(self, parser):
        super(ImportRattail, self).add_parser_args(parser)
        parser.add_argument('--dbkey', metavar='KEY', default='host',
                            help="Config key for database engine to be used as the Rattail "
                            "\"host\", i.e. the source of the data to be imported.  This key "
                            "must be defined in the [rattail.db] section of your config file.  "
                            "Defaults to 'host'.")

    def get_handler_kwargs(self, **kwargs):
        if 'args' in kwargs:
            kwargs['dbkey'] = kwargs['args'].dbkey
        return kwargs


class ImportRattailBulk(ImportRattail):
    """
    Bulk-import data from another Rattail database
    """
    name = 'import-rattail-bulk'
    description = __doc__.strip()
    handler_key = 'rattail_bulk'
    default_handler_spec = 'rattail.importing.rattail_bulk:BulkFromRattailToRattail'