Coverage for pgsql_upserter / temp_staging.py: 89.01%

1"""Temporary table management for PostgreSQL upsert operations.""" 

2 

3import json 

4import logging 

5import uuid 

6import psycopg2 

7 

8from psycopg2.extras import RealDictCursor, execute_values 

9from typing import Any 

10 

11from .exceptions import PgsqlUpserterError 

12 

13logger = logging.getLogger(__name__) 

14 

15 

16def _normalize_null_values(value: Any) -> Any | None: 

17 """Convert common null representations to None for PostgreSQL NULL.""" 

18 if value is None: 

19 return None 

20 

21 if isinstance(value, str): 

22 normalized = value.strip().lower() 

23 if normalized in ['', 'none', 'null', 'nan', 'na', '-']: 

24 return None 

25 

26 return value 

27 

28 

29def _convert_value_for_postgres(value: Any, column_data_type: str) -> Any | None: 

30 """Convert Python values to PostgreSQL-compatible formats based on column data type. 

31 

32 Args: 

33 value: The value to convert 

34 column_data_type: PostgreSQL data type of the target column 

35 

36 Returns: 

37 Converted value suitable for PostgreSQL insertion 

38 """ 

39 # First apply null normalization 

40 normalized_value = _normalize_null_values(value) 

41 if normalized_value is None: 

42 return None 

43 

44 # Convert based on PostgreSQL data type 

45 column_type = column_data_type.lower() 

46 

47 try: 

48 # Handle JSONB and JSON columns 

49 if column_type in ('jsonb', 'json'): 

50 if isinstance(normalized_value, (dict, list)): 

51 return json.dumps(normalized_value) 

52 elif isinstance(normalized_value, str): 52 ↛ 60line 52 didn't jump to line 60 because the condition on line 52 was always true

53 # Try to parse as JSON, if it fails, treat as string literal 

54 try: 

55 json.loads(normalized_value) 

56 return normalized_value # Already valid JSON string 

57 except json.JSONDecodeError: 

58 return json.dumps(normalized_value) # Wrap non-JSON string in quotes 

59 else: 

60 return json.dumps(normalized_value) 

61 

62 # Handle PostgreSQL arrays (text[], integer[], etc.) 

63 elif column_type.endswith('[]') or column_type.startswith('_'): 

64 if isinstance(normalized_value, list): 

65 # Convert Python list to PostgreSQL array format 

66 # For text arrays, we need to escape quotes and handle null values 

67 array_elements = [] 

68 for item in normalized_value: 

69 if item is None: 

70 array_elements.append('NULL') 

71 elif isinstance(item, str): 71 ↛ 76line 71 didn't jump to line 76 because the condition on line 71 was always true

72 # Escape quotes and wrap in quotes 

73 escaped = item.replace('"', '\\"') 

74 array_elements.append(f'"{escaped}"') 

75 else: 

76 array_elements.append(str(item)) 

77 return '{' + ','.join(array_elements) + '}' 

78 elif isinstance(normalized_value, str): 78 ↛ 83line 78 didn't jump to line 83 because the condition on line 78 was always true

79 # If it's already a string, assume it's in PostgreSQL array format 

80 return normalized_value 

81 else: 

82 # Convert single value to single-element array 

83 if isinstance(normalized_value, str): 

84 escaped = normalized_value.replace('"', '\\"') 

85 return '{' + f'"{escaped}"' + '}' 

86 else: 

87 return '{' + str(normalized_value) + '}' 

88 

89 # For all other types, return the normalized value as-is 

90 # PostgreSQL driver will handle the conversion 

91 else: 

92 return normalized_value 

93 

94 except Exception as e: 

95 logger.warning(f"Failed to convert value {repr(value)} for column type {column_data_type}: {e}") 

96 # Fallback: convert to string 

97 return str(normalized_value) if normalized_value is not None else None 

98 

99 

100def create_temp_table(connection, target_table: str, schema: str = 'public') -> str: 

101 """Create temporary table with same structure as target table. 

102 

103 Args: 

104 connection: Active PostgreSQL connection 

105 target_table: Name of the target table to copy structure from 

106 schema: Schema name (default: 'public') 

107 

108 Returns: 

109 str: The temporary table name that was created 

110 

111 Raises: 

112 PgsqlUpserterError: If temp table creation fails 

113 """ 

114 # Generate unique temp table name 

115 temp_table_name = f"temp_staging_{uuid.uuid4().hex[:8]}" 

116 

117 try: 

118 with connection.cursor(cursor_factory=RealDictCursor) as cursor: 

119 # Create temp table with same structure, no constraints or defaults 

120 create_sql = f""" 

121 CREATE TEMPORARY TABLE {temp_table_name} 

122 (LIKE {schema}.{target_table} EXCLUDING ALL) 

123 """ 

124 

125 cursor.execute(create_sql) 

126 

127 # Drop auto-generated columns to avoid constraint issues 

128 # Get auto-generated column info 

129 cursor.execute(""" 

130 SELECT column_name, is_generated, column_default, data_type 

131 FROM information_schema.columns 

132 WHERE table_schema = %s AND table_name = %s 

133 ORDER BY ordinal_position 

134 """, (schema, target_table)) 

135 

136 columns_to_drop = [] 

137 for row in cursor.fetchall(): 

138 # Check if column is auto-generated 

139 is_auto_generated = ( 

140 row['data_type'] in ('serial', 'bigserial') or 

141 (row['column_default'] and 'nextval(' in row['column_default'].lower()) or 

142 row['is_generated'] == 'ALWAYS' or 

143 (row['column_default'] and any(ts in row['column_default'].lower() 

144 for ts in ['current_timestamp', 'now()', 'clock_timestamp()'])) 

145 ) 

146 

147 if is_auto_generated: 

148 columns_to_drop.append(row['column_name']) 

149 

150 # Drop auto-generated columns from temp table 

151 for col_name in columns_to_drop: 

152 cursor.execute(f"ALTER TABLE {temp_table_name} DROP COLUMN {col_name}") 

153 

154 connection.commit() 

155 

156 logger.debug( 

157 f"Created temporary table '{temp_table_name}' based on '{schema}.{target_table}', dropped {len(columns_to_drop)} auto-generated columns") # noqa 

158 return temp_table_name 

159 

160 except psycopg2.Error as e: 

161 connection.rollback() 

162 raise PgsqlUpserterError(f"Failed to create temporary table: {e}") 

163 

164 

165def bulk_insert_to_temp( 

166 connection, 

167 temp_table_name: str, 

168 data_list: list[dict[str, Any]], 

169 matched_columns: list[str], 

170 target_schema=None, 

171 batch_size: int = 1000, 

172 show_progress: bool = True 

173) -> int: 

174 """Bulk insert filtered data into temporary table. 

175 

176 Args: 

177 connection: Active PostgreSQL connection 

178 temp_table_name: Name of the temporary table 

179 data_list: List of dictionaries containing data to insert 

180 matched_columns: List of column names to include in insert 

181 target_schema: TableSchema object for data type conversion (optional) 

182 batch_size: Number of rows to process in each batch 

183 show_progress: Whether to show progress for large datasets 

184 

185 Returns: 

186 int: Number of rows inserted 

187 

188 Raises: 

189 PgsqlUpserterError: If bulk insert fails 

190 """ 

191 if not data_list or not matched_columns: 

192 return 0 

193 

194 total_rows = len(data_list) 

195 if show_progress and total_rows > batch_size: 

196 logger.info(f"Processing {total_rows} rows...") 

197 

198 # Build column data type mapping for proper conversion 

199 column_type_map = {} 

200 if target_schema: 

201 for col_info in target_schema.columns: 

202 if col_info.name in matched_columns: 

203 column_type_map[col_info.name] = col_info.data_type 

204 

205 try: 

206 with connection.cursor(cursor_factory=RealDictCursor) as cursor: 

207 # Filter and normalize data with progress tracking 

208 filtered_data = [] 

209 for i, row in enumerate(data_list): 

210 filtered_row = [] 

211 for col in matched_columns: 

212 value = row.get(col) # Missing keys become None 

213 

214 # Convert value based on column data type if available 

215 if col in column_type_map: 

216 converted_value = _convert_value_for_postgres(value, column_type_map[col]) 

217 else: 

218 converted_value = _normalize_null_values(value) 

219 

220 filtered_row.append(converted_value) 

221 filtered_data.append(filtered_row) 

222 

223 # Show progress every batch_size rows 

224 if show_progress and (i + 1) % batch_size == 0: 

225 progress = (i + 1) / total_rows * 100 

226 logger.info(f"Processed {i + 1}/{total_rows} rows ({progress:.1f}%)") 

227 

228 # Build INSERT statement for execute_values 

229 columns_sql = ', '.join(matched_columns) 

230 insert_sql = f"INSERT INTO {temp_table_name} ({columns_sql}) VALUES %s" 

231 

232 if show_progress: 

233 logger.info("Inserting data into temporary table...") 

234 

235 # Use execute_values for better performance in serverless environments 

236 execute_values( 

237 cursor, 

238 insert_sql, 

239 filtered_data, 

240 template=None, 

241 page_size=batch_size # Good balance for serverless memory limits 

242 ) 

243 

244 rows_inserted = len(filtered_data) # Use actual data length instead of cursor.rowcount 

245 connection.commit() 

246 

247 if show_progress: 

248 logger.info(f"Successfully inserted {rows_inserted} batched rows") 

249 

250 logger.info(f"Inserted {rows_inserted} total rows into temporary table '{temp_table_name}'") 

251 return rows_inserted 

252 

253 except psycopg2.Error as e: 

254 connection.rollback() 

255 # Try to cleanup temp table 

256 _cleanup_temp_table(connection, temp_table_name) 

257 raise PgsqlUpserterError(f"Failed to bulk insert into temporary table: {e}") 

258 

259 

260# Alias for better API naming 

261populate_temp_table = bulk_insert_to_temp 

262 

263 

264def convert_temp_to_permanent( 

265 connection, 

266 temp_table_name: str, 

267 permanent_table_name: str, 

268 schema: str = 'public' 

269) -> None: 

270 """Convert temporary table to permanent table for debugging purposes. 

271 

272 Args: 

273 connection: Active PostgreSQL connection 

274 temp_table_name: Name of the temporary table 

275 permanent_table_name: Name for the permanent table (will be dropped if exists) 

276 schema: Schema name (default: 'public') 

277 

278 Raises: 

279 PgsqlUpserterError: If conversion fails 

280 """ 

281 try: 

282 with connection.cursor() as cursor: 

283 # Drop existing permanent table if it exists 

284 cursor.execute(f"DROP TABLE IF EXISTS {schema}.{permanent_table_name}") 

285 

286 # Create permanent table from temp table structure and data 

287 cursor.execute(f""" 

288 CREATE TABLE {schema}.{permanent_table_name} AS 

289 SELECT * FROM {temp_table_name} 

290 """) 

291 

292 connection.commit() 

293 

294 # Get row count for confirmation 

295 cursor.execute(f"SELECT COUNT(*) FROM {schema}.{permanent_table_name}") 

296 row_count = cursor.fetchone()[0] 

297 

298 logger.info( 

299 f"Converted temp table '{temp_table_name}' to permanent table '{schema}.{permanent_table_name}' with {row_count} rows") # noqa 

300 

301 except psycopg2.Error as e: 

302 connection.rollback() 

303 raise PgsqlUpserterError(f"Failed to convert temp table to permanent: {e}") 

304 

305 

306def _cleanup_temp_table(connection, temp_table_name: str) -> None: 

307 """Silently attempt to drop temporary table.""" 

308 try: 

309 with connection.cursor() as cursor: 

310 cursor.execute(f"DROP TABLE IF EXISTS {temp_table_name}") 

311 connection.commit() 

312 logger.debug(f"Cleaned up temporary table '{temp_table_name}'") 

313 except Exception: 

314 # Fail silently as requested - temp tables auto-cleanup on session end anyway 

315 pass