diff --git a/swh/storage/in_memory.py b/swh/storage/in_memory.py --- a/swh/storage/in_memory.py +++ b/swh/storage/in_memory.py @@ -1277,22 +1277,16 @@ for visit in visits: if isinstance(visit['date'], str): visit['date'] = dateutil.parser.parse(visit['date']) - if isinstance(visit['origin'], str): - origin = \ - self.origin_get([{'url': visit['origin']}])[0] - if not origin: - raise ValueError('Unknown origin: %s' % visit['origin']) - visit['origin'] = origin['id'] + visit['origin'] = self.origin_get([visit['origin']])[0] if self.journal_writer: for visit in visits: - visit = visit.copy() - visit['origin'] = self.origin_get([{'id': visit['origin']}])[0] + visit = copy.deepcopy(visit) del visit['origin']['id'] self.journal_writer.write_addition('origin_visit', visit) for visit in visits: - origin_id = visit['origin'] + origin_id = visit['origin']['id'] visit_id = visit['visit'] self._objects[(origin_id, visit_id)].append( diff --git a/swh/storage/storage.py b/swh/storage/storage.py --- a/swh/storage/storage.py +++ b/swh/storage/storage.py @@ -1287,22 +1287,19 @@ for visit in visits: if isinstance(visit['date'], str): visit['date'] = dateutil.parser.parse(visit['date']) - if isinstance(visit['origin'], str): - visit['origin'] = \ - self.origin_get({'url': visit['origin']})['id'] + visit['origin'] = \ + self.origin_get([visit['origin']], db=db, cur=cur)[0] if self.journal_writer: for visit in visits: - visit = visit.copy() - origin = self.origin_get( - [{'id': visit['origin']}], db=db, cur=cur)[0] - visit['origin'] = origin + visit = copy.deepcopy(visit) if visit.get('type') is None: - visit['type'] = origin['type'] + visit['type'] = visit['origin']['type'] del visit['origin']['id'] self.journal_writer.write_addition('origin_visit', visit) for visit in visits: + visit['origin'] = visit['origin']['id'] # TODO: upsert them all in a single query db.origin_visit_upsert(**visit, cur=cur) diff --git a/swh/storage/tests/test_storage.py b/swh/storage/tests/test_storage.py --- a/swh/storage/tests/test_storage.py +++ b/swh/storage/tests/test_storage.py @@ -1959,7 +1959,7 @@ # when self.storage.origin_visit_upsert([ { - 'origin': origin_id_or_url, + 'origin': self.origin2, 'date': self.date_visit2, 'visit': 123, 'type': self.origin2['type'], @@ -1968,7 +1968,7 @@ 'snapshot': None, }, { - 'origin': origin_id_or_url, + 'origin': self.origin2, 'date': '2018-01-01 23:00:00+00', 'visit': 1234, 'type': self.origin2['type'], @@ -2042,7 +2042,7 @@ origin_id_or_url, date=self.date_visit2) self.storage.origin_visit_upsert([{ - 'origin': origin_id_or_url, + 'origin': self.origin2, 'date': self.date_visit2, 'visit': origin_visit1['visit'], 'type': self.origin2['type'],