Python postgres_copy.CopyMapping类代码示例

 def test_field_value_mapping_and_types(self):
     c = CopyMapping(
         BasicMockObject,
         self.mapping_path,
         dict(name='NAME', number='NUMBER', dt='DATE'),
         field_value_mapping={
             'name': {
                 'ben': 'Master Ben',
                 'joe': 'Padawan Joe',
                 'jane': 'Jedi Jane'
             },
             'number': {
                 'seven': 7,
                 'three': 3,
                 'five': 5
             }
         },
         field_copy_types={'number': 'text'}
     )
     c.save()
     self.assertEqual(BasicMockObject.objects.count(), 3)
     self.assertEqual(
         list(BasicMockObject.objects.order_by('name').values_list('name', 'number')),
         [('Jedi Jane', 5), ('Master Ben', 7), ('Padawan Joe', 3)]
     )
     self.assertEqual(
         BasicMockObject.objects.get(name='Master Ben').dt,
         date(2012, 1, 1)
     )

    def handle(self, *args, **options):
        data = os.path.join(settings.BASE_DIR, "data")
        files = list(all_files(data, "*.csv"))

        for filepath in files:
            log("Opening file {}\n".format(filepath), "cyan")

            log("  Loading data ...\n")

            copy = CopyMapping(
                Incident,
                filepath,
                dict(
                    date="IncidentDate",
                    year="Year",
                    raw_location="Location",
                    address="Edited Street Address",
                    city_and_state="City and State",
                    full_address="Edited Full Address",
                    latitude="Latitude",
                    longitude="Longitude",
                    accuracy_score="Accuracy Score",
                    accuracy_type="Accuracy Type",
                    number="Number",
                    street="Street",
                    city="City",
                    state="State",
                    county="County",
                    zipcode="Zip",
                ),
            )

            copy.save()

            log("    Data loaded!\n", "green")

 def test_bad_static_values(self):
     with self.assertRaises(ValueError):
         c = CopyMapping(
             ExtendedMockObject,
             self.name_path,
             dict(name='NAME', number='NUMBER', dt='DATE'),
             encoding='UTF-8',
             static_mapping={'static_bad':1,}
         )
         c.save()

 def handle(self, *args, **kwargs):
     c = CopyMapping(
         # Give it the model
         Person,
         # The path to your CSV
         '/path/to/my/data.csv',
         # And a dict mapping the  model fields to CSV headers
         dict(name='NAME', number='NUMBER', dt='DATE')
     )
     # Then save it.
     c.save()

 def test_limited_save(self):
     c = CopyMapping(
         LimitedMockObject,
         self.name_path,
         dict(name='NAME', dt='DATE')
     )
     c.save()
     self.assertEqual(LimitedMockObject.objects.count(), 3)
     self.assertEqual(
         LimitedMockObject.objects.get(name='BEN').dt,
         date(2012, 1, 1)
     )

 def test_field_override_save(self):
     c = CopyMapping(
         MockObject,
         self.null_path,
         dict(name='NAME', number='NUMBER', dt='DATE'),
     )
     c.save()
     self.assertEqual(MockObject.objects.count(), 5)
     self.assertEqual(MockObject.objects.get(name='BADBOY').number, None)
     self.assertEqual(
         MockObject.objects.get(name='BEN').dt,
         date(2012, 1, 1)
     )

 def test_silent_save(self):
     c = CopyMapping(
         MockObject,
         self.name_path,
         dict(name='NAME', number='NUMBER', dt='DATE'),
     )
     c.save(silent=True)
     self.assertEqual(MockObject.objects.count(), 3)
     self.assertEqual(MockObject.objects.get(name='BEN').number, 1)
     self.assertEqual(
         MockObject.objects.get(name='BEN').dt,
         date(2012, 1, 1)
     )

 def test_pipe_save(self):
     c = CopyMapping(
         MockObject,
         self.pipe_path,
         dict(name='NAME', number='NUMBER', dt='DATE'),
         delimiter="|",
     )
     c.save()
     self.assertEqual(MockObject.objects.count(), 3)
     self.assertEqual(MockObject.objects.get(name='BEN').number, 1)
     self.assertEqual(
         MockObject.objects.get(name='BEN').dt,
         date(2012, 1, 1)
     )

    def test_save_foreign_key(self):
        c = CopyMapping(
            MockObject,
            self.foreign_path,
            dict(name='NAME', number='NUMBER', dt='DATE', parent='PARENT')
        )

        c.save()
        self.assertEqual(MockObject.objects.count(), 3)
        self.assertEqual(MockObject.objects.get(name='BEN').parent_id, 4)
        self.assertEqual(
            MockObject.objects.get(name='BEN').dt,
            date(2012, 1, 1)
        )

 def test_static_values(self):
     c = CopyMapping(
         ExtendedMockObject,
         self.name_path,
         dict(name='NAME', number='NUMBER', dt='DATE'),
         static_mapping={'static_val':1,'static_string':'test'}
     )
     c.save()
     self.assertEqual(
         ExtendedMockObject.objects.filter(static_val = 1).count(),
         3
     )
     self.assertEqual(
         ExtendedMockObject.objects.filter(static_string = 'test').count(),
         3
     )

 def test_overload_save(self):
     c = CopyMapping(
         OverloadMockObject,
         self.name_path,
         dict(name='NAME', number='NUMBER', dt='DATE'),
         overloaded_mapping=dict(lower_name='NAME')
     )
     c.save()
     self.assertEqual(OverloadMockObject.objects.count(), 3)
     self.assertEqual(OverloadMockObject.objects.get(name='BEN').number, 1)
     self.assertEqual(OverloadMockObject.objects.get(lower_name='ben').number, 1)
     self.assertEqual(
         OverloadMockObject.objects.get(name='BEN').dt,
         date(2012, 1, 1)
     )
     omo = OverloadMockObject.objects.first()
     self.assertEqual(omo.name.lower(), omo.lower_name)

    def load_postgresql(self):
        """
        Load the file into a PostgreSQL database using COPY
        """
        # Drop all the records from the target model's real table
        self.cursor.execute('TRUNCATE TABLE "%s" CASCADE' % (self.model._meta.db_table))

        c = CopyMapping(
            self.model, self.csv, dict((f.name, f.db_column) for f in self.model._meta.fields), using=self.database
        )
        c.save(silent=True)

        # Print out the results
        if self.verbosity > 2:
            csv_count = self.get_row_count()
            model_count = self.model.objects.count()
            self.finish_load_message(model_count, csv_count)

    def load_postgresql(self, model, csv_path):
        """
        Takes a model and a csv_path and loads it into postgresql
        """
        # Drop all the records from the target model's real table
        self.cursor.execute('TRUNCATE TABLE "%s" CASCADE' % (
            model._meta.db_table
        ))

        c = CopyMapping(
            model,
            csv_path,
            dict((f.name, f.db_column) for f in model._meta.fields),
        )
        c.save(silent=True)

        # Print out the results
        csv_count = self.get_row_count(csv_path)
        model_count = model.objects.count()
        self.finish_load_message(model_count, csv_count)

    def test_foreignkey_mapping(self):
        c = CopyMapping(
            BasicMockObject,
            self.name_path,
            dict(name='NAME', number='NUMBER', dt='DATE')
        )
        c.save()
        self.assertEqual(BasicMockObject.objects.count(), 3)

        c = CopyMapping(
            BasicMockObject,
            self.fk_mapping_path,
            dict(name='NAME', number='NUMBER', dt='DATE', parent='PARENT'),
            field_value_mapping={
                'parent': {
                    'ben': BasicMockObject.objects.get(name='ben').pk,
                    'joe': BasicMockObject.objects.get(name='joe').pk,
                    'jane': BasicMockObject.objects.get(name='jane').pk
                },
            },
            field_copy_types={'parent': 'text'}
        )
        c.save()
        self.assertEqual(BasicMockObject.objects.count(), 6)
        self.assertEqual(
            list(BasicMockObject.objects.order_by('name').
                 values_list('name', 'parent__name')),
            [('ben', None), ('ben junior', 'ben'), ('jane', None),
             ('jane junior', 'jane'), ('joe', None), ('joe junior', 'joe')]
        )

    def import_csv(self, request):
        if 'csv' not in request.FILES:
            return JsonResponse(
                status=status.HTTP_400_BAD_REQUEST,
                data={
                    'success': False,
                    'errors': ["Missing 'csv' parameter"]
                }
            )

        csv = request.FILES['csv']

        file_path = "/tmp/{0}.csv".format(str(uuid.uuid4()))

        with open(file_path, 'wb+') as destination:
            for chunk in csv.chunks():
                destination.write(chunk)

        try:
            copy_mapping = CopyMapping(
                models.Concept,
                file_path,
                ConceptViewSet.CSV_COLUMN_MAPPING
            )
            copy_mapping.save()
        except (ValueError, DatabaseError):
            expected_columns = ', '.join(ConceptViewSet.CSV_COLUMN_MAPPING.keys())
            return JsonResponse(
                status=status.HTTP_400_BAD_REQUEST,
                data={
                    'success': False,
                    'errors': ['CSV import error. Expected columns are {0}'.format(expected_columns)]
                }
            )
        finally:
            os.remove(file_path)

        return JsonResponse({
            'success': True,
        })