Как ускорить запись в базу данных?

Question

У меня есть функция, которая ищет файлы json в каталоге, анализирует файл и записывает данные в базу данных. Моя проблема - запись в базу данных, потому что это занимает около 30 минут. Любая идея, как я могу ускорить запись в базу данных? У меня есть несколько довольно больших файлов для анализа, но разбор файла не проблема, это занимает около 3 минут. В настоящее время я использую sqlite, но в будущем я изменю его на PostgreSQL.

Вот моя функция:

def create_database():
    with transaction.atomic():
        directory = os.fsencode('data/web_files/unzip')
        for file in os.listdir(directory):
            filename = os.fsdecode(file)

            with open('data/web_files/unzip/{}'.format(filename.strip()), encoding="utf8") as f:
                data = json.load(f)
                cve_items = data['CVE_Items']
                for i in range(len(cve_items)):
                    database_object = DataNist()

                    try:
                        impact = cve_items[i]['impact']['baseMetricV2']
                        database_object.severity = impact['severity']
                        database_object.exp_score = impact['exploitabilityScore']
                        database_object.impact_score = impact['impactScore']
                        database_object.cvss_score = impact['cvssV2']['baseScore']

                    except KeyError:
                        database_object.severity = ''
                        database_object.exp_score = ''
                        database_object.impact_score = ''
                        database_object.cvss_score = ''

                    for vendor_data in cve_items[i]['cve']['affects']['vendor']['vendor_data']:
                            database_object.vendor_name = vendor_data['vendor_name']

                            for description_data in cve_items[i]['cve']['description']['description_data']:
                                database_object.description = description_data['value']

                            for product_data in vendor_data['product']['product_data']:
                                database_object.product_name = product_data['product_name']
                                database_object.save()

                                for version_data in product_data['version']['version_data']:
                                    if version_data['version_value'] != '-':
                                    database_object.versions_set.create(version=version_data['version_value'])

My models.py:

class DataNist(models.Model):
    vendor_name = models.CharField(max_length=100)
    product_name = models.CharField(max_length=100)
    description = models.TextField()
    date = models.DateTimeField(default=timezone.now)
    severity = models.CharField(max_length=10)
    exp_score = models.IntegerField()
    impact_score = models.IntegerField()
    cvss_score = models.IntegerField()


    def __str__(self):
        return self.vendor_name + "-" + self.product_name


class Versions(models.Model):
    data = models.ForeignKey(DataNist, on_delete=models.CASCADE)
    version = models.CharField(max_length=50)

    def __str__(self):
        return self.version

Буду признателен, если вы дадите мне какой-нибудь совет, как мне улучшить мой код.

Answer 1

Хорошо, учитывая структуру данных, что-то вроде этого может работать для вас.

Это автономный код, кроме этого вызова .objects.bulk_create();как прокомментировано в коде, два определенных класса на самом деле будут моделями в вашем приложении Django.(Кстати, вы, вероятно, также хотите сохранить идентификатор CVE как уникальное поле.)

В вашем исходном коде было ошибочное предположение, что у каждой "конечной записи" в данных затронутой версии будет один и тот же поставщик, которыйможет не быть правдой.Вот почему структура модели здесь имеет отдельную модель продукта-версии, которая имеет поля поставщика, продукта и версии.(Если вы хотите немного оптимизировать вещи, вы можете дедуплицировать AffectedProductVersion s даже через DataNist s (что, помимо прочего, не является идеальным названием для модели)).

ИКонечно, как вы уже сделали в своем исходном коде, импорт должен выполняться внутри транзакции (transaction.atomic()).

Надеюсь, это поможет.

import json
import os
import types


class DataNist(types.SimpleNamespace):  # this would actually be a model
    severity = ""
    exp_score = ""
    impact_score = ""
    cvss_score = ""

    def save(self):
        pass


class AffectedProductVersion(types.SimpleNamespace):  # this too
    # (foreign key to DataNist here)
    vendor_name = ""
    product_name = ""
    version_value = ""


def import_item(item):
    database_object = DataNist()
    try:
        impact = item["impact"]["baseMetricV2"]
    except KeyError:  # no impact object available
        pass
    else:
        database_object.severity = impact.get("severity", "")
        database_object.exp_score = impact.get("exploitabilityScore", "")
        database_object.impact_score = impact.get("impactScore", "")
        if "cvssV2" in impact:
            database_object.cvss_score = impact["cvssV2"]["baseScore"]

    for description_data in item["cve"]["description"]["description_data"]:
        database_object.description = description_data["value"]
        break  # only grab the first description

    database_object.save()  # save the base object

    affected_versions = []
    for vendor_data in item["cve"]["affects"]["vendor"]["vendor_data"]:
        for product_data in vendor_data["product"]["product_data"]:
            for version_data in product_data["version"]["version_data"]:
                affected_versions.append(
                    AffectedProductVersion(
                        data_nist=database_object,
                        vendor_name=vendor_data["vendor_name"],
                        product_name=product_data["product_name"],
                        version_name=version_data["version_value"],
                    )
                )

    AffectedProductVersion.objects.bulk_create(
        affected_versions
    )  # save all the version information

    return database_object  # in case the caller needs it


with open("nvdcve-1.0-2019.json") as infp:
    data = json.load(infp)
    for item in data["CVE_Items"]:
        import_item(item)