Coverage for utilities/forms/bulk_import.py: 18%

1import csv 

2import json 

3from io import StringIO 

4 

5import yaml 

6from django import forms 

7from django.utils.translation import gettext as _ 

8 

9from core.forms.mixins import SyncedDataMixin 

10from netbox.choices import CSVDelimiterChoices, ImportFormatChoices, ImportMethodChoices 

11from netbox.forms.mixins import ChangelogMessageMixin 

12from utilities.constants import CSV_DELIMITERS 

13from utilities.forms.mixins import BackgroundJobMixin 

14from utilities.forms.utils import parse_csv 

15 

16 

17class BulkImportForm(ChangelogMessageMixin, BackgroundJobMixin, SyncedDataMixin, forms.Form): 

18 import_method = forms.ChoiceField( 

19 choices=ImportMethodChoices, 

20 required=False 

21 ) 

22 data = forms.CharField( 

23 required=False, 

24 widget=forms.Textarea(attrs={'class': 'font-monospace'}), 

25 help_text=_("Enter object data in CSV, JSON or YAML format.") 

26 ) 

27 upload_file = forms.FileField( 

28 label=_("Data file"), 

29 required=False 

30 ) 

31 format = forms.ChoiceField( 

32 choices=ImportFormatChoices, 

33 initial=ImportFormatChoices.AUTO 

34 ) 

35 csv_delimiter = forms.ChoiceField( 

36 choices=CSVDelimiterChoices, 

37 initial=CSVDelimiterChoices.AUTO, 

38 label=_("CSV delimiter"), 

39 help_text=_("The character which delimits CSV fields. Applies only to CSV format."), 

40 required=False 

41 ) 

42 

43 data_field = 'data' 

44 

45 def clean(self): 

46 super().clean() 

47 

48 # Determine import method 

49 import_method = self.cleaned_data.get('import_method') or ImportMethodChoices.DIRECT 

50 

51 # Determine whether we're reading from form data or an uploaded file 

52 if self.cleaned_data['data'] and import_method != ImportMethodChoices.DIRECT: 

53 raise forms.ValidationError(_("Form data must be empty when uploading/selecting a file.")) 

54 if import_method == ImportMethodChoices.UPLOAD: 

55 self.data_field = 'upload_file' 

56 file = self.files.get('upload_file') 

57 data = file.read().decode('utf-8-sig') 

58 elif import_method == ImportMethodChoices.DATA_FILE: 

59 self.data_field = 'data_file' 

60 data = self.cleaned_data['data_file'].data_as_string 

61 else: 

62 data = self.cleaned_data['data'] 

63 

64 # Determine the data format 

65 if self.cleaned_data['format'] == ImportFormatChoices.AUTO: 

66 if self.cleaned_data['csv_delimiter'] != CSVDelimiterChoices.AUTO: 

67 # Specifying the CSV delimiter implies CSV format 

68 format = ImportFormatChoices.CSV 

69 else: 

70 format = self._detect_format(data) 

71 else: 

72 format = self.cleaned_data['format'] 

73 

74 # Process data according to the selected format 

75 if format == ImportFormatChoices.CSV: 

76 delimiter = self.cleaned_data.get('csv_delimiter', CSVDelimiterChoices.AUTO) 

77 self.cleaned_data['data'] = self._clean_csv(data, delimiter=delimiter) 

78 elif format == ImportFormatChoices.JSON: 

79 self.cleaned_data['data'] = self._clean_json(data) 

80 elif format == ImportFormatChoices.YAML: 

81 self.cleaned_data['data'] = self._clean_yaml(data) 

82 else: 

83 raise forms.ValidationError(_("Unknown data format: {format}").format(format=format)) 

84 

85 def _detect_format(self, data): 

86 """ 

87 Attempt to automatically detect the format (CSV, JSON, or YAML) of the given data, or raise 

88 a ValidationError. 

89 """ 

90 try: 

91 if data[0] in ('{', '['): 

92 return ImportFormatChoices.JSON 

93 if data.startswith('---') or data.startswith('- '): 

94 return ImportFormatChoices.YAML 

95 # Look for any of the CSV delimiters in the first line (ignoring the default 'auto' choice) 

96 first_line = data.split('\n', 1)[0] 

97 csv_delimiters = CSV_DELIMITERS.values() 

98 if any(x in first_line for x in csv_delimiters): 

99 return ImportFormatChoices.CSV 

100 except IndexError: 

101 pass 

102 raise forms.ValidationError({ 

103 'format': _('Unable to detect data format. Please specify.') 

104 }) 

105 

106 def _clean_csv(self, data, delimiter=CSVDelimiterChoices.AUTO): 

107 """ 

108 Clean CSV-formatted data. The first row will be treated as column headers. 

109 """ 

110 # Determine the CSV dialect 

111 if delimiter == CSVDelimiterChoices.AUTO: 

112 # This uses a rough heuristic to detect the CSV dialect based on the presence of supported delimiting 

113 # characters. If the data is malformed, we'll fall back to the default Excel dialect. 

114 delimiters = ''.join(CSV_DELIMITERS.values()) 

115 try: 

116 dialect = csv.Sniffer().sniff(data.strip(), delimiters=delimiters) 

117 except csv.Error: 

118 dialect = csv.excel 

119 elif delimiter in (CSVDelimiterChoices.COMMA, CSVDelimiterChoices.SEMICOLON, CSVDelimiterChoices.PIPE): 

120 dialect = csv.excel 

121 dialect.delimiter = delimiter 

122 elif delimiter == CSVDelimiterChoices.TAB: 

123 dialect = csv.excel_tab 

124 else: 

125 raise forms.ValidationError({ 

126 'csv_delimiter': _('Invalid CSV delimiter'), 

127 }) 

128 

129 stream = StringIO(data.strip()) 

130 reader = csv.reader(stream, dialect=dialect) 

131 headers, records = parse_csv(reader) 

132 

133 # Set CSV headers for reference by the model form 

134 headers.pop('id', None) 

135 self._csv_headers = headers 

136 

137 return records 

138 

139 def _clean_json(self, data): 

140 """ 

141 Clean JSON-formatted data. If only a single object is defined, it will be encapsulated as a list. 

142 """ 

143 try: 

144 data = json.loads(data) 

145 # Accommodate for users entering single objects 

146 if type(data) is not list: 

147 data = [data] 

148 return data 

149 except json.decoder.JSONDecodeError as err: 

150 raise forms.ValidationError({ 

151 self.data_field: f"Invalid JSON data: {err}" 

152 }) 

153 

154 def _clean_yaml(self, data): 

155 """ 

156 Clean YAML-formatted data. Data must be either 

157 a) A single document comprising a list of dictionaries (each representing an object), or 

158 b) Multiple documents, separated with the '---' token 

159 """ 

160 records = [] 

161 try: 

162 for data in yaml.load_all(data, Loader=yaml.SafeLoader): 

163 if type(data) is list: 

164 records.extend(data) 

165 elif type(data) is dict: 

166 records.append(data) 

167 else: 

168 raise forms.ValidationError({ 

169 self.data_field: _( 

170 "Invalid YAML data. Data must be in the form of multiple documents, or a single document " 

171 "comprising a list of dictionaries." 

172 ) 

173 }) 

174 except yaml.error.YAMLError as err: 

175 raise forms.ValidationError({ 

176 self.data_field: f"Invalid YAML data: {err}" 

177 }) 

178 

179 return records