Skip to content

Use get or create to reduce the number of objects created #58

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
singh1114 wants to merge 1 commit into from
Closed

Use get or create to reduce the number of objects created #58

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 22 additions & 20 deletions vulnerabilities/data_dump.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,39 +34,39 @@ def debian_dump(extract_data):
Save data scraped from Debian' security tracker.
"""
for data in extract_data:
vulnerability = Vulnerability.objects.create(
vulnerability, _ = Vulnerability.objects.get_or_create(
summary=data.get('description', ''),
)
VulnerabilityReference.objects.create(
VulnerabilityReference.objects.get_or_create(
vulnerability=vulnerability,
reference_id=data.get('vulnerability_id', ''),
)

pkg_name = data.get('package_name', '')
package = Package.objects.create(
package, _ = Package.objects.get_or_create(
name=pkg_name,
version=data.get('version', ''),
)

if data['status'] == 'open':
ImpactedPackage.objects.create(
ImpactedPackage.objects.get_or_create(
vulnerability=vulnerability,
package=package
)
else:
ResolvedPackage.objects.create(
ResolvedPackage.objects.get_or_create(
vulnerability=vulnerability,
package=package
)

fixed_version = data.get('fixed_version')
if fixed_version:
package = Package.objects.create(
package, _ = Package.objects.get_or_create(
name=pkg_name,
version=fixed_version,
)

ResolvedPackage.objects.create(
ResolvedPackage.objects.get_or_create(
vulnerability=vulnerability,
package=package
)
Expand All @@ -77,17 +77,19 @@ def ubuntu_dump(html):
Dump data scraped from Ubuntu's security tracker.
"""
for data in html:
vulnerability = Vulnerability.objects.create(
summary='',
# If we don't do this, every VulnerabilityReference will have single
# Vulnerability.
vulnerability, _ = Vulnerability.objects.get_or_create(
summary='{} vulnerability'.format(data.get('cve_id')),
)
VulnerabilityReference.objects.create(
VulnerabilityReference.objects.get_or_create(
vulnerability=vulnerability,
reference_id=data.get('cve_id'),
)
package = Package.objects.create(
package, _ = Package.objects.get_or_create(
name=data.get('package_name'),
)
ImpactedPackage.objects.create(
ImpactedPackage.objects.get_or_create(
vulnerability=vulnerability,
package=package
)
Expand All @@ -111,39 +113,39 @@ def archlinux_dump(extract_data):
if not fixed_version:
fixed_version = 'None'

vulnerability = Vulnerability.objects.create(
vulnerability, _ = Vulnerability.objects.get_or_create(
summary=item['type'],
)

for vulnerability_id in vulnerabilities:
VulnerabilityReference.objects.create(
VulnerabilityReference.objects.get_or_create(
vulnerability=vulnerability,
reference_id=vulnerability_id,
url='https://security.archlinux.org/{}'.format(vulnerability_id)
)

for package_name in packages_name:
package_affected = Package.objects.create(
package_affected, _ = Package.objects.get_or_create(
name=package_name,
version=affected_version
)
ImpactedPackage.objects.create(
ImpactedPackage.objects.get_or_create(
vulnerability=vulnerability,
package=package_affected
)
PackageReference.objects.create(
PackageReference.objects.get_or_create(
package=package_affected,
repository='https://security.archlinux.org/package/{}'.format(package_name)
)
package_fixed = Package.objects.create(
package_fixed, _ = Package.objects.get_or_create(
name=package_name,
version=fixed_version
)
ResolvedPackage.objects.create(
ResolvedPackage.objects.get_or_create(
vulnerability=vulnerability,
package=package_fixed
)
PackageReference.objects.create(
PackageReference.objects.get_or_create(
package=package_fixed,
repository='https://security.archlinux.org/package/{}'.format(package_name)
)
6 changes: 3 additions & 3 deletions vulnerabilities/tests/test_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ def test_debian_response(self):
debian_dump(extract_data)
response = self.client.get('/api/packages/?name=mimetex', format='json').data

self.assertEqual(4, response['count'])
self.assertEqual(2, response['count'])

first_result = response['results'][0]
self.assertEqual('mimetex', first_result['name'])
Expand All @@ -69,7 +69,7 @@ def test_ubuntu_response(self):
"version": "",
"platform": "",
"vulnerabilities": [{
"summary": "",
"summary": "CVE-2012-3386 vulnerability",
"cvss": None,
"references": [{
"reference_id": "CVE-2012-3386",
Expand All @@ -93,7 +93,7 @@ def test_serializers(self):
pk = Package.objects.filter(name="mimetex")
response = PackageSerializer(pk, many=True).data

self.assertEqual(4, len(response))
self.assertEqual(2, len(response))

first_result = response[0]
self.assertEqual('mimetex', first_result['name'])
Expand Down
73 changes: 64 additions & 9 deletions vulnerabilities/tests/test_data_dump.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,8 @@ def setUpTestData(self):
with open(os.path.join(TEST_DATA, 'debian.json')) as f:
test_data = json.load(f)

extract_data = debian.extract_vulnerabilities(test_data)
debian_dump(extract_data)
self.extract_data = debian.extract_vulnerabilities(test_data)
debian_dump(self.extract_data)

def test_Vulnerability(self):
"""
Expand Down Expand Up @@ -83,9 +83,7 @@ def test_Package(self):
"""
Check that all packages from the test data are stored in the database
"""
# There are five rows in Package because currently the models allow duplicates
# (see issue #28).
self.assertEqual(5, Package.objects.count())
self.assertEqual(3, Package.objects.count())

self.assertTrue(Package.objects.filter(name='mimetex'))
self.assertTrue(Package.objects.get(name='librsync'))
Expand Down Expand Up @@ -114,15 +112,33 @@ def test_ResolvedPackage(self):
self.assertIn('1.50-1.1', versions)
self.assertIn('1.74-1', versions)

def test_debian_data_dump_twice(self):
"""
Scrape data from Debian' main tracker, save it
in the database and verify entries.
"""
debian_dump(self.extract_data)

self.assertEqual(3, Vulnerability.objects.count())
self.assertEqual(3, VulnerabilityReference.objects.count())
self.assertEqual(3, Package.objects.count())

# Dumping the data twice doesn't create new objects.
debian_dump(self.extract_data)

self.assertEqual(3, Vulnerability.objects.count())
self.assertEqual(3, VulnerabilityReference.objects.count())
self.assertEqual(3, Package.objects.count())


class TestUbuntuDataDump(TestCase):
@classmethod
def setUpTestData(self):
with open(os.path.join(TEST_DATA, 'ubuntu_main.html')) as f:
test_data = f.read()

data = ubuntu.extract_cves(test_data)
ubuntu_dump(data)
self.data = ubuntu.extract_cves(test_data)
ubuntu_dump(self.data)

def test_data_dump(self):
"""
Expand All @@ -132,15 +148,32 @@ def test_data_dump(self):
self.assertEqual(reference.reference_id, 'CVE-2002-2439')
self.assertTrue(Package.objects.filter(name='gcc-4.6')[0].name, 'gcc-4.6')

def test_ubuntu_data_dump_twice(self):
"""
Scrape data from Ubuntu twice from main tracker, save it
in the database and verify single time entry.
"""
ubuntu_dump(self.data)
count = Package.objects.all().count()
reference = VulnerabilityReference.objects.filter(
reference_id='CVE-2002-2439')
self.assertEqual(reference[0].reference_id, 'CVE-2002-2439')
self.assertTrue(Package.objects.filter(name='gcc-4.6')[0].name,
'gcc-4.6')

# Dumping the data twice doesn't create new objects.
ubuntu_dump(self.data)
self.assertEqual(count, Package.objects.all().count())


class TestArchLinuxDataDump(TestCase):

@classmethod
def setUpTestData(self):
with open(os.path.join(TEST_DATA, 'archlinux.json')) as f:
test_data = json.load(f)
self.test_data = json.load(f)

archlinux_dump(test_data)
archlinux_dump(self.test_data)

def test_Vulnerability(self):
"""
Expand Down Expand Up @@ -190,3 +223,25 @@ def test_ResolvedPackage(self):

self.assertEqual(4, len(resolved_pkgs))
self.assertEqual('2.6.1-1', resolved_pkg.package.version)

def test_archlinux_data_dump_twice(self):
"""
Scrape data from Archlinux' main tracker twice, save it
in the database and verify no multiple entries.
"""
archlinux_dump(self.test_data)
self.assertEqual(1, Vulnerability.objects.count())
self.assertEqual(14, VulnerabilityReference.objects.count())
self.assertEqual(8, Package.objects.count())
self.assertEqual(8, PackageReference.objects.count())
self.assertEqual(4, ImpactedPackage.objects.count())
self.assertEqual(4, ResolvedPackage.objects.count())

# Dumping the data twice doesn't create new objects.
archlinux_dump(self.test_data)
self.assertEqual(1, Vulnerability.objects.count())
self.assertEqual(14, VulnerabilityReference.objects.count())
self.assertEqual(8, Package.objects.count())
self.assertEqual(8, PackageReference.objects.count())
self.assertEqual(4, ImpactedPackage.objects.count())
self.assertEqual(4, ResolvedPackage.objects.count())