Completed UI changes to download V2 JSON (#1204)

Author: brian-kim31

app/airflow/dags/libs/rules_export/core.py

@@ -40,13 +40,22 @@ def pre_process_rules(**kwargs) -> None:
     validated_params = pull_validated_params(kwargs, "validate_params_rules_export")
     scan_report_id = validated_params["scan_report_id"]
     file_type = validated_params["file_type"]
+
+    # Map file types to temp table suffixes
+    if file_type in ["application/json_v1", "application/json_v2"]:
+        temp_table_suffix = "json"
+    elif file_type == "csv":
+        temp_table_suffix = "csv"
+    else:
+        temp_table_suffix = file_type
+
     try:
         #  Create or update the temp table with the processed rules data
         pg_hook.run(
             create_update_temp_rules_table_query
             % {
                 "scan_report_id": scan_report_id,
-                "file_type": file_type,
+                "file_type": temp_table_suffix,
             },
         )
     except Exception as e:
@@ -70,6 +79,9 @@ def build_and_upload_rules_file(**kwargs) -> None:
     user_id = validated_params["user_id"]
     scan_report_name = validated_params["scan_report_name"]
     file_type = validated_params["file_type"]
+    json_version = validated_params.get(
+        "json_version", "v1"
+    )  # default to v1 if not specified
 
     try:
         # Setup file config. (Credit: @AndyRae)
@@ -79,13 +91,27 @@ def build_and_upload_rules_file(**kwargs) -> None:
                 "mapping_csv",
                 "csv",
             ),
+            "application/json_v1": FileHandlerConfig(
+                lambda: build_rules_json(scan_report_name, scan_report_id),
+                "mapping_json",
+                "json",
+            ),
+            "application/json_v2": FileHandlerConfig(
+                lambda: build_rules_json_v2(scan_report_name, scan_report_id),
+                "mapping_json_v2",
+                "json",
+            ),
             "json": FileHandlerConfig(
                 lambda: (
                     build_rules_json_v2(scan_report_name, scan_report_id)
                     if AIRFLOW_VAR_JSON_VERSION == "v2"
                     else build_rules_json(scan_report_name, scan_report_id)
                 ),
-                "mapping_json",
+                (
+                    "mapping_json_v2"
+                    if AIRFLOW_VAR_JSON_VERSION == "v2"
+                    else "mapping_json"
+                ),
                 "json",
             ),
         }
@@ -97,7 +123,11 @@ def build_and_upload_rules_file(**kwargs) -> None:
         file_extension = config.file_extension
 
         # build file name
-        filename = f"Rules - {scan_report_name} - {scan_report_id} - {datetime.now()}.{file_extension}"
+        if file_type in ["application/json_v1", "application/json_v2"]:
+            version = "V1" if file_type == "application/json_v1" else "V2"
+            filename = f"Rules - {scan_report_name} - {scan_report_id} - {version} - {datetime.now()}.{file_extension}"
+        else:
+            filename = f"Rules - {scan_report_name} - {scan_report_id} - {datetime.now()}.{file_extension}"
 
         # Upload to blob storage
         upload_blob_to_storage(
@@ -146,7 +176,11 @@ def build_and_upload_rules_file(**kwargs) -> None:
                 "DROP TABLE IF EXISTS temp_rules_export_%(scan_report_id)s_%(file_type)s"
                 % {
                     "scan_report_id": scan_report_id,
-                    "file_type": file_type,
+                    "file_type": (
+                        "json"
+                        if file_type in ["application/json_v1", "application/json_v2"]
+                        else file_type
+                    ),
                 }
             )
         except Exception as e:
@@ -159,7 +193,7 @@ def build_and_upload_rules_file(**kwargs) -> None:
             )
             raise e
     except Exception as e:
-        logging.error(f"Error creating file entry: {str(e)}")
+        logging.error(f"Error building and uploading rules file: {str(e)}")
         update_job_status(
             scan_report=scan_report_id,
             stage=JobStageType.DOWNLOAD_RULES,

app/airflow/dags/libs/utils.py

@@ -270,12 +270,19 @@ def _validate_dag_params(
             if param == "file_type":
                 if value == "application/json":
                     value = "json"
+                elif value == "application/json_v1":
+                    value = "application/json_v1"
+                elif value == "application/json_v2":
+                    value = "application/json_v2"
                 elif value == "text/csv":
                     value = "csv"
                 else:
                     errors.append(
-                        f"Invalid {param}: {value}. Must be application/json or text/csv."
+                        f"Invalid {param}: {value}. Must be application/json, application/json_v1, application/json_v2, or text/csv."
                     )
+            elif param == "json_version":
+                if value not in ["v1", "v2"]:
+                    errors.append(f"Invalid {param}: {value}. Must be v1 or v2.")
             validated_params[param] = value
 
     # Validate boolean parameters
@@ -357,7 +364,7 @@ def validate_params_SR_processing(**context):
 def validate_params_rules_export(**context):
     """Validates parameters required for rules export DAG tasks."""
     int_params = ["scan_report_id", "user_id"]
-    string_params = ["file_type", "scan_report_name"]
+    string_params = ["file_type", "scan_report_name", "json_version"]
     return _validate_dag_params(
         int_params=int_params,
         string_params=string_params,

app/api/files/fixtures/filetypes.json

@@ -7,6 +7,7 @@
           "display_name": "Mapping Rules JSON"
       }
   },
+  
   {
       "model": "files.filetype",
       "pk": 2,
@@ -38,5 +39,13 @@
           "value": "data_dictionary",
           "display_name": "Data Dictionary"
       }
+  },
+  {
+    "model": "files.filetype",
+    "pk": 6,
+    "fields": {
+        "value": "mapping_json_v2",
+        "display_name": "Mapping Rules JSON V2"
+    }
   }
 ]

app/api/files/views.py

@@ -77,7 +77,7 @@ def post(self, request, *args, **kwargs):
             - scan_report_id (int): The ID of the scan report for which the
             file is to be generated.
             - file_type (str): The type of file to generate (e.g.,
-            'application/json' or 'text/csv').
+            'application/json_v1', 'application/json_v2', or 'text/csv').
 
         Upon successful validation of the input, a message is sent to the
         Rules Export Queue, and a job record is created in the database to
@@ -111,20 +111,35 @@ def post(self, request, *args, **kwargs):
                 )
             # Get the scan report model to get the scan report name for both Azure and Airflow tasks later on
             scan_report = ScanReport.objects.get(id=scan_report_id)
+
+            # Determine the JSON version for the message
+            json_version = "v1"  # default
+            if file_type == "application/json_v2":
+                json_version = "v2"
+            elif file_type == "application/json_v1":
+                json_version = "v1"
+
             msg = {
                 "scan_report_id": scan_report_id,
                 "scan_report_name": scan_report.dataset,
                 "user_id": request.user.id,
                 "file_type": file_type,
+                "json_version": json_version,
             }
 
             worker_service.trigger_rules_export(msg)
+
             # Create job record for downloading file
+            file_type_description = (
+                "JSON V1"
+                if file_type == "application/json_v1"
+                else "JSON V2" if file_type == "application/json_v2" else "CSV"
+            )
             Job.objects.create(
                 scan_report=ScanReport.objects.get(id=scan_report_id),
                 stage=JobStage.objects.get(value="DOWNLOAD_RULES"),
                 status=StageStatus.objects.get(value="IN_PROGRESS"),
-                details=f"A Mapping Rules {'JSON' if file_type == 'application/json' else 'CSV'} is being generated.",
+                details=f"A Mapping Rules {file_type_description} is being generated.",
             )
         except json.JSONDecodeError:
             return JsonResponse({"error": "Invalid JSON data."}, status=400)

app/next-client-app/app/(protected)/scanreports/[id]/actions-download-menu.tsx

@@ -15,7 +15,9 @@ type Props = { scanreportId: string };
 export function ActionsDownloadMenu({ scanreportId }: Props) {
   const router = useRouter();
 
-  const handleDownload = async (fileType: FileTypeFormat) => {
+  const handleDownload = async (
+    fileType: FileTypeFormat | "application/json_v1" | "application/json_v2"
+  ) => {
     const resp = await requestFile(Number(scanreportId), fileType);
     if (resp.success) {
       router.push(`/scanreports/${scanreportId}/downloads`);
@@ -30,10 +32,14 @@ export function ActionsDownloadMenu({ scanreportId }: Props) {
   return (
     <DropdownMenuGroup>
       <DropdownMenuLabel>Downloads</DropdownMenuLabel>
-      <DropdownMenuItem onSelect={() => handleDownload("application/json")}>
+      <DropdownMenuItem onSelect={() => handleDownload("application/json_v1")}>
         <FileJson />
         Mapping JSON
       </DropdownMenuItem>
+      <DropdownMenuItem onSelect={() => handleDownload("application/json_v2")}>
+        <FileJson />
+        Mapping JSON V2
+      </DropdownMenuItem>
       <DropdownMenuItem onSelect={() => handleDownload("text/csv")}>
         <FileSpreadsheet />
         Mapping CSV

app/next-client-app/app/(protected)/scanreports/[id]/downloads/columns.tsx

@@ -26,7 +26,7 @@ export const columns: ColumnDef<FileDownload>[] = [
       return format(created_at, "d MMM HH:mm");
     },
     enableHiding: true,
-    enableSorting: true,
+    enableSorting: true
   },
   {
     id: "User",
@@ -43,7 +43,7 @@ export const columns: ColumnDef<FileDownload>[] = [
       return <>{user.username}</>;
     },
     enableHiding: true,
-    enableSorting: false,
+    enableSorting: false
   },
   {
     id: "Type",
@@ -60,7 +60,7 @@ export const columns: ColumnDef<FileDownload>[] = [
       return <Badge variant="outline">{file_type.display_name}</Badge>;
     },
     enableHiding: true,
-    enableSorting: false,
+    enableSorting: false
   },
   {
     id: "Download",
@@ -85,6 +85,6 @@ export const columns: ColumnDef<FileDownload>[] = [
       );
     },
     enableHiding: true,
-    enableSorting: false,
-  },
+    enableSorting: false
+  }
 ];

app/next-client-app/components/recommendations/stored-recommendations-button.tsx

@@ -111,7 +111,7 @@ export function StoredRecommendationsButton({
             >
               <Sparkles className="h-4 w-4 text-purple-500" />
               {mappingRecommendations && mappingRecommendations.length > 0
-                ? `Recommendations (${mappingRecommendations.length})`
+                ? "Recommendations"
                 : "No Recommendations"}
             </Button>
           </div>

app/next-client-app/types/files.ts

@@ -1,15 +1,17 @@
-type FileTypeFormat = "application/json" | "image/svg+xml" | "text/csv";
+
+type FileTypeFormat = "application/json" | "application/json_v1" | "application/json_v2" | "image/svg+xml" | "text/csv";
 type FileTypeValue =
   | "mapping_json"
+  | "mapping_json_v2"
   | "mapping_csv"
   | "mapping_svg"
   | "data_dictionary"
   | "scan_report";
-
 interface FileType {
   value: FileTypeValue;
   display_name:
     | "Mapping Rules JSON"
+    | "Mapping Rules JSON V2"
     | "Mapping Rules CSV"
     | "Mapping Rules SVG"
     | "Data Dictionary"