diff --git a/rattail/commands/__init__.py b/rattail/commands/__init__.py index c4a2d32c..d5560837 100644 --- a/rattail/commands/__init__.py +++ b/rattail/commands/__init__.py @@ -2,7 +2,7 @@ ################################################################################ # # Rattail -- Retail Software Framework -# Copyright © 2010-2017 Lance Edgar +# Copyright © 2010-2018 Lance Edgar # # This file is part of Rattail. # @@ -27,4 +27,4 @@ Console Commands from __future__ import unicode_literals, absolute_import from .core import main, Command, Subcommand, Dump, date_argument, list_argument -from .importing import ImportSubcommand +from .importing import ImportSubcommand, ImportFromCSV diff --git a/rattail/commands/importing.py b/rattail/commands/importing.py index ce3c9f8d..32b61ba9 100644 --- a/rattail/commands/importing.py +++ b/rattail/commands/importing.py @@ -241,6 +241,18 @@ class ImportSubcommand(Subcommand): log.info("transaction was committed") +class ImportFromCSV(ImportSubcommand): + """ + Generic base class for commands which import from a CSV file. + """ + + def add_parser_args(self, parser): + super(ImportFromCSV, self).add_parser_args(parser) + + parser.add_argument('--source-csv', metavar='PATH', required=True, + help="Path to CSV file to be used as data source.") + + class ExportRattail(ImportSubcommand): """ Export data to another Rattail database diff --git a/rattail/importing/__init__.py b/rattail/importing/__init__.py index 313c68fd..860e0e4a 100644 --- a/rattail/importing/__init__.py +++ b/rattail/importing/__init__.py @@ -1,8 +1,8 @@ -# -*- coding: utf-8 -*- +# -*- coding: utf-8; -*- ################################################################################ # # Rattail -- Retail Software Framework -# Copyright © 2010-2017 Lance Edgar +# Copyright © 2010-2018 Lance Edgar # # This file is part of Rattail. # @@ -26,7 +26,7 @@ Data Importing Framework from __future__ import unicode_literals, absolute_import -from .importers import Importer, FromQuery, FromDjango, BatchImporter, BulkImporter +from .importers import Importer, FromCSV, FromQuery, FromDjango, BatchImporter, BulkImporter from .sqlalchemy import FromSQLAlchemy, ToSQLAlchemy from .postgresql import BulkToPostgreSQL from .handlers import ImportHandler, BulkImportHandler, FromSQLAlchemyHandler, ToSQLAlchemyHandler diff --git a/rattail/importing/importers.py b/rattail/importing/importers.py index f7eeeeb8..a1f1a095 100644 --- a/rattail/importing/importers.py +++ b/rattail/importing/importers.py @@ -34,6 +34,7 @@ from rattail.db import cache from rattail.db.util import QuerySequence from rattail.time import make_utc from rattail.util import data_diffs, progress_loop, OrderedDict +from rattail.csvutil import UnicodeDictReader log = logging.getLogger(__name__) @@ -144,6 +145,11 @@ class Importer(object): Perform any setup necessary, e.g. cache lookups for existing data. """ + def datasync_setup(self): + """ + Perform any setup necessary, in the context of a datasync job. + """ + def teardown(self): """ Perform any cleanup after import, if necessary. @@ -563,6 +569,29 @@ class Importer(object): data[field], data[field2] = data[field2], None +class FromCSV(Importer): + """ + Generic base class for importers whose data source is a CSV file. + """ + + def setup(self): + if not hasattr(self, 'source_csv_path'): + if hasattr(self, 'args') and hasattr(self.args, 'source_csv'): + self.source_csv_path = self.args.source_csv + + def get_host_objects(self): + if six.PY3: + source_csv_file = open(self.source_csv_path, 'rt', encoding='latin_1') + reader = UnicodeDictReader(source_csv_file) + else: + source_csv_file = open(self.source_csv_path, 'rb') + reader = UnicodeDictReader(source_csv_file, 'latin_1') + + objects = list(reader) + source_csv_file.close() + return objects + + class FromQuery(Importer): """ Generic base class for importers whose raw external data source is a