@@ -44,7 +44,7 @@ def pre_process_rules(**kwargs) -> None:
4444 scan_report = scan_report_id ,
4545 stage = JobStageType .DOWNLOAD_RULES ,
4646 status = StageStatusType .FAILED ,
47- details = f"Error creating or updating the temp table for rules export for scan report { scan_report_id } ,
47+ details = f"Error creating or updating the temp table: { str ( e ) } ,
4848 )
4949 raise e
5050
@@ -60,85 +60,95 @@ def build_and_upload_rules_file(**kwargs) -> None:
6060 scan_report_name = validated_params ["scan_report_name" ]
6161 file_type = validated_params ["file_type" ]
6262
63- # Setup file config. (Credit: @AndyRae)
64- file_handlers : Dict [str , FileHandlerConfig ] = {
65- "csv" : FileHandlerConfig (
66- lambda : build_rules_csv (scan_report_id ),
67- "mapping_csv" ,
68- "csv" ,
69- ),
70- "json" : FileHandlerConfig (
71- lambda : build_rules_json (scan_report_name , scan_report_id ),
72- "mapping_json" ,
73- "json" ,
74- ),
75- }
76- config = file_handlers [file_type ]
77-
78- # Generate it
79- file = config .handler ()
80- file_type_value = config .file_type_value
81- file_extension = config .file_extension
63+ try :
64+ # Setup file config. (Credit: @AndyRae)
65+ file_handlers : Dict [str , FileHandlerConfig ] = {
66+ "csv" : FileHandlerConfig (
67+ lambda : build_rules_csv (scan_report_id ),
68+ "mapping_csv" ,
69+ "csv" ,
70+ ),
71+ "json" : FileHandlerConfig (
72+ lambda : build_rules_json (scan_report_name , scan_report_id ),
73+ "mapping_json" ,
74+ "json" ,
75+ ),
76+ }
77+ config = file_handlers [file_type ]
8278
83- # build file name
84- filename = f"Rules - { scan_report_name } { scan_report_id } { datetime .now ()} { file_extension }
79+ # Generate it
80+ file = config .handler ()
81+ file_type_value = config .file_type_value
82+ file_extension = config .file_extension
8583
86- # Upload to blob storage
87- upload_blob_to_storage (
88- container_name = "rules-exports" ,
89- blob_name = filename ,
90- data = file ,
91- content_type = file_type ,
92- )
84+ # build file name
85+ filename = f"Rules - { scan_report_name } { scan_report_id } { datetime .now ()} { file_extension }
9386
94- # Get file_type_id from files_filetype table
95- file_type_id_query = """
96- SELECT id FROM files_filetype WHERE value = %(file_type_value)s
97- """
98- file_type_id_result = pg_hook .get_first (
99- file_type_id_query , parameters = {"file_type_value" : file_type_value }
100- )
101- if not file_type_id_result :
102- raise ValueError (f"No file type found for value: { file_type_value } )
103- file_type_id = file_type_id_result [0 ]
104-
105- try :
106- # Create file entry in the database
107- pg_hook .run (
108- create_file_entry_query ,
109- parameters = {
110- "name" : filename ,
111- "file_url" : filename ,
112- "user_id" : user_id ,
113- "file_type_id" : file_type_id ,
114- "scan_report_id" : scan_report_id ,
115- "created_at" : datetime .now (),
116- "updated_at" : datetime .now (),
117- },
118- )
119- update_job_status (
120- scan_report = scan_report_id ,
121- stage = JobStageType .DOWNLOAD_RULES ,
122- status = StageStatusType .COMPLETE ,
123- details = "Rules file successfully created and uploaded to blob storage" ,
87+ # Upload to blob storage
88+ upload_blob_to_storage (
89+ container_name = "rules-exports" ,
90+ blob_name = filename ,
91+ data = file ,
92+ content_type = file_type ,
12493 )
12594
126- if AIRFLOW_DEBUG_MODE == "true" :
127- return
128- # Clean up
129- pg_hook .run (
130- "DROP TABLE IF EXISTS temp_rules_export_%(scan_report_id)s_%(file_type)s"
131- % {
132- "scan_report_id" : scan_report_id ,
133- "file_type" : file_type ,
134- }
95+ # Get file_type_id from files_filetype table
96+ file_type_id_query = """
97+ SELECT id FROM files_filetype WHERE value = %(file_type_value)s
98+ """
99+ file_type_id_result = pg_hook .get_first (
100+ file_type_id_query , parameters = {"file_type_value" : file_type_value }
135101 )
102+ if not file_type_id_result :
103+ raise ValueError (f"No file type found for value: { file_type_value } )
104+ file_type_id = file_type_id_result [0 ]
105+
106+ try :
107+ # Create file entry in the database
108+ pg_hook .run (
109+ create_file_entry_query ,
110+ parameters = {
111+ "name" : filename ,
112+ "file_url" : filename ,
113+ "user_id" : user_id ,
114+ "file_type_id" : file_type_id ,
115+ "scan_report_id" : scan_report_id ,
116+ "created_at" : datetime .now (),
117+ "updated_at" : datetime .now (),
118+ },
119+ )
120+ update_job_status (
121+ scan_report = scan_report_id ,
122+ stage = JobStageType .DOWNLOAD_RULES ,
123+ status = StageStatusType .COMPLETE ,
124+ details = "Rules file successfully created and uploaded to blob storage" ,
125+ )
126+
127+ if AIRFLOW_DEBUG_MODE == "true" :
128+ return
129+ # Clean up
130+ pg_hook .run (
131+ "DROP TABLE IF EXISTS temp_rules_export_%(scan_report_id)s_%(file_type)s"
132+ % {
133+ "scan_report_id" : scan_report_id ,
134+ "file_type" : file_type ,
135+ }
136+ )
137+ except Exception as e :
138+ logging .error (f"Error creating file entry: { str (e )} )
139+ update_job_status (
140+ scan_report = scan_report_id ,
141+ stage = JobStageType .DOWNLOAD_RULES ,
142+ status = StageStatusType .FAILED ,
143+ details = f"Export rules file failed: { str (e )} ,
144+ )
145+ raise e
136146 except Exception as e :
137147 logging .error (f"Error creating file entry: { str (e )} )
138148 update_job_status (
139149 scan_report = scan_report_id ,
140150 stage = JobStageType .DOWNLOAD_RULES ,
141151 status = StageStatusType .FAILED ,
142- details = f"Error creating file entry in the database for scan report { scan_report_id } ,
152+ details = f"Export rules file failed: { str ( e ) } ,
143153 )
144154 raise e
0 commit comments