Skip to content

Commit 8905f62

Browse files
ziadhanyziadhany
committed
Try to improve export performance
Signed-off-by: ziadhany <ziadhany2016@gmail.com>
1 parent e58f279 commit 8905f62

2 files changed

Lines changed: 49 additions & 72 deletions

File tree

‎vulnerabilities/management/commands/export.py‎

Lines changed: 49 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -38,11 +38,51 @@ def export_data(self, git_path):
3838
by run `python manage.py export /path/vulnerablecode-data`
3939
"""
4040
self.stdout.write(f"Exporting vulnerablecode data")
41-
41+
old_purl_without_version = None
42+
old_package_data_with_header = None
43+
old_pkg_filepath = None
4244
for purl in Package.objects.all().paginated():
45+
purl_without_version = PackageURL(
46+
type=purl.type,
47+
namespace=purl.namespace,
48+
name=purl.name,
49+
)
4350
package_dir = create_sub_paths(git_path, purl.type, purl.namespace, purl.name)
44-
pck_filepath = package_dir.joinpath(f"{purl.type}-{purl.name}.yaml")
45-
write_package_data(pck_filepath, purl)
51+
if old_purl_without_version and (
52+
purl_without_version.type == old_purl_without_version.type
53+
and purl_without_version.namespace == old_purl_without_version.namespace
54+
and purl_without_version.name == old_purl_without_version.name
55+
):
56+
package_data = {
57+
"purl": str(purl),
58+
"affected_by_vulnerabilities": [
59+
vuln.vulnerability_id for vuln in purl.affected_by
60+
],
61+
"fixing_vulnerabilities": [vuln.vulnerability_id for vuln in purl.fixing],
62+
}
63+
old_package_data_with_header["versions"].append(package_data)
64+
else:
65+
if old_pkg_filepath and old_purl_without_version:
66+
with open(old_pkg_filepath, mode="w") as f:
67+
yaml_data = saneyaml.dump(old_package_data_with_header)
68+
f.write(yaml_data)
69+
70+
old_pkg_filepath = package_dir.joinpath(f"{purl.type}-{purl.name}.yaml")
71+
old_purl_without_version = purl_without_version
72+
old_package_data_with_header = {
73+
"pacakge": str(purl_without_version),
74+
"versions": [
75+
{
76+
"purl": str(purl),
77+
"affected_by_vulnerabilities": [
78+
vuln.vulnerability_id for vuln in purl.affected_by
79+
],
80+
"fixing_vulnerabilities": [
81+
vuln.vulnerability_id for vuln in purl.fixing
82+
],
83+
}
84+
],
85+
}
4686

4787
for vul in purl.vulnerabilities.all():
4888
vul_filepath = package_dir.joinpath(f"{vul.vulnerability_id}.yaml")
@@ -72,53 +112,6 @@ def write_vul_data(filepath, vul):
72112
f.write(vul_data)
73113

74114

75-
def write_package_data(filepath, purl):
76-
"""
77-
write the pacakge data in a file with a path like this
78-
`path/ecosystem/package/package_type-package_name` and create the directories if it doesn't exist
79-
and if the pacakge file doesn't exist we create the new header
80-
"""
81-
purl_object = PackageURL.from_string(purl.package_url)
82-
purl_without_version = PackageURL(
83-
type=purl_object.type,
84-
namespace=purl_object.namespace,
85-
name=purl_object.name,
86-
)
87-
88-
package_data_with_header = saneyaml.dump(
89-
{
90-
"pacakge": str(purl_without_version),
91-
"versions": [
92-
{
93-
"purl": str(purl_object),
94-
"affected_by_vulnerabilities": [
95-
vuln.vulnerability_id for vuln in purl.affected_by
96-
],
97-
"fixing_vulnerabilities": [vuln.vulnerability_id for vuln in purl.fixing],
98-
}
99-
],
100-
}
101-
)
102-
103-
package_data = {
104-
"purl": str(purl_object),
105-
"affected_by_vulnerabilities": [vuln.vulnerability_id for vuln in purl.affected_by],
106-
"fixing_vulnerabilities": [vuln.vulnerability_id for vuln in purl.fixing],
107-
}
108-
109-
if Path(filepath).is_file():
110-
with open(filepath, "r") as f:
111-
old_yaml = saneyaml.load(f)
112-
old_yaml["versions"].append(package_data)
113-
114-
if old_yaml:
115-
with open(filepath, "w") as f:
116-
f.write(saneyaml.dump(old_yaml))
117-
else:
118-
with open(filepath, encoding="utf-8", mode="w") as f:
119-
f.write(package_data_with_header)
120-
121-
122115
def create_sub_paths(git_path, purl_type, purl_namespace, purl_name):
123116
"""
124117
create the directories if it doesn't exist : `path/purl_type/purl_namespace/purl_name`
@@ -127,9 +120,12 @@ def create_sub_paths(git_path, purl_type, purl_namespace, purl_name):
127120
if not ecosystem_dir.is_dir():
128121
ecosystem_dir.mkdir()
129122

130-
namespace_dir = ecosystem_dir.joinpath(purl_namespace)
131-
if not namespace_dir.is_dir():
132-
namespace_dir.mkdir()
123+
if purl_namespace:
124+
namespace_dir = ecosystem_dir.joinpath(purl_namespace)
125+
if not namespace_dir.is_dir():
126+
namespace_dir.mkdir()
127+
else:
128+
namespace_dir = ecosystem_dir
133129

134130
package_dir = namespace_dir.joinpath(purl_name)
135131
if not package_dir.is_dir():

‎vulnerabilities/tests/test_export.py‎

Lines changed: 0 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
from django.core.management.base import CommandError
1010

1111
from vulnerabilities.management.commands.export import create_sub_paths
12-
from vulnerabilities.management.commands.export import write_package_data
1312
from vulnerabilities.management.commands.export import write_vul_data
1413
from vulnerabilities.models import Package
1514
from vulnerabilities.models import PackageRelatedVulnerability
@@ -78,24 +77,6 @@ def test_write_vul_data(tmp_path, vulnerability):
7877
assert Path(pck_filepath).read_text() == saneyaml.dump(expected_data)
7978

8079

81-
@pytest.mark.django_db
82-
def test_write_package_data(tmp_path, package_related_vulnerability):
83-
expected_data = {
84-
"pacakge": "pkg:generic/nginx/test",
85-
"versions": [
86-
{
87-
"purl": "pkg:generic/nginx/test@2",
88-
"affected_by_vulnerabilities": ["VCID-pst6-b358-aaap"],
89-
"fixing_vulnerabilities": [],
90-
}
91-
],
92-
}
93-
pck_filepath = os.path.join(tmp_path, "filename")
94-
write_package_data(pck_filepath, package_related_vulnerability)
95-
assert os.path.isfile(pck_filepath)
96-
assert Path(pck_filepath).read_text() == saneyaml.dump(expected_data)
97-
98-
9980
class TestExportCommand(TestCase):
10081
def test_missing_path(self):
10182
with pytest.raises(CommandError) as cm:

0 commit comments

Comments
 (0)