Add migration to load default catalogs to DB

CHANGELOG.md

@@ -9,6 +9,11 @@ v0.9.10-dev (August xx, 2021)
 * Add the catalog_key to statement's `sid_class` and `source` fields when adding new statement to a component in library.
 
 
+**Data fix**
+
+* Add migration in controls to load default control catalogs into CatalogData in database.
+
+
 v0.9.9 (August 12, 2021)
 ------------------------
 

controls/migrations/0060_auto_20210816_1634.py

@@ -0,0 +1,51 @@
+# Generated by Django 3.2.5 on 2021-08-16 16:34
+
+from django.db import migrations
+import os.path
+import json
+
+from controls.models import Element
+from controls.oscal import CatalogData
+
+
+def load_catalog_data(apps, schema_editor):
+    """Load control catalog data into database"""
+
+    # Load the default control catalogs and baselines
+    CATALOG_PATH = os.path.join(os.path.dirname(__file__),'..','data','catalogs')
+    BASELINE_PATH = os.path.join(os.path.dirname(__file__),'..','data','baselines')
+
+    # TODO: Check directory exists
+    catalog_files = [file for file in os.listdir(CATALOG_PATH) if file.endswith('.json')]
+    # Load catalog and baseline data into database records from source files if data records do not exist in database
+    for cf in catalog_files:
+        catalog_key = cf.replace("_catalog.json", "")
+        with open(os.path.join(CATALOG_PATH,cf), 'r') as json_file:
+            catalog_json = json.load(json_file)
+        baseline_filename = cf.replace("_catalog.json", "_baselines.json")
+        if os.path.isfile(os.path.join(BASELINE_PATH, baseline_filename)):
+            with open(os.path.join(BASELINE_PATH, baseline_filename), 'r') as json_file:
+                baselines_json = json.load(json_file)
+        else:
+            baselines_json = {}
+
+        catalog, created = CatalogData.objects.get_or_create(
+                catalog_key=catalog_key,
+                catalog_json=catalog_json,
+                baselines_json=baselines_json
+            )
+        if created:
+            print(f"{catalog_key} record created into database")
+        else:
+                print(f"{catalog_key} record found in database")
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('controls', '0059_auto_20210811_0001'),
+    ]
+
+    operations = [
+         migrations.RunPython(load_catalog_data)
+    ]