Coverage for pgsql_upserter / conflict_resolver.py: 80.88%

1"""Conflict resolution logic for PostgreSQL upsert operations.""" 

2 

3import logging 

4 

5from dataclasses import dataclass 

6 

7from .schema_inspector import TableSchema 

8from .exceptions import PgsqlUpserterError 

9 

10# Configure module logger 

11logger = logging.getLogger(__name__) 

12 

13 

14@dataclass 

15class ConflictStrategy: 

16 """Represents a conflict resolution strategy.""" 

17 type: str # "PRIMARY_KEY", "UNIQUE_COMBINED", "INSERT_ONLY" 

18 columns: list[str] # Column names for conflict resolution 

19 description: str # Human-readable description 

20 

21 

22@dataclass 

23class DeduplicationResult: 

24 """Results from temp table deduplication process.""" 

25 original_count: int 

26 deduplicated_count: int 

27 dropped_count: int 

28 dropped_reasons: dict[str, int] # reason -> count mapping 

29 

30 

31def find_conflict_strategy( 

32 table_schema: TableSchema, 

33 matched_columns: list[str] 

34) -> ConflictStrategy: 

35 """ 

36 Determine the best conflict resolution strategy based on table schema and available columns. 

37 

38 Priority: 

39 1. Non-auto-generated primary key(s) - highest priority 

40 2. All unique constraints combined - medium priority 

41 3. Simple insert only - fallback 

42 

43 Args: 

44 table_schema: Table schema information 

45 matched_columns: List of columns available in the data (from column_matcher) 

46 

47 Returns: 

48 ConflictStrategy object with strategy details 

49 

50 Raises: 

51 PgsqlUpserterError: If no valid strategy can be determined 

52 """ 

53 logger.debug(f"Analyzing conflict strategy for table '{table_schema.table_name}' " 

54 f"with {len(matched_columns)} available columns") 

55 

56 # Convert matched_columns to set for faster lookup 

57 available_columns = set(matched_columns) 

58 

59 # Strategy 1: Check for non-auto-generated primary key 

60 if table_schema.primary_key: 60 ↛ 78line 60 didn't jump to line 78 because the condition on line 60 was always true

61 pk_columns = table_schema.primary_key.columns 

62 # Check if all PK columns are available in matched_columns 

63 pk_available = all(col in available_columns for col in pk_columns) 

64 

65 if pk_available: 65 ↛ 66line 65 didn't jump to line 66 because the condition on line 65 was never true

66 strategy = ConflictStrategy( 

67 type="PRIMARY_KEY", 

68 columns=pk_columns, 

69 description=f"Primary key conflict resolution on: {pk_columns}" 

70 ) 

71 logger.debug(f"Using PRIMARY_KEY strategy: {strategy.description}") 

72 return strategy 

73 else: 

74 missing_pk_cols = [col for col in pk_columns if col not in available_columns] 

75 logger.warning(f"Primary key columns not available in data: {missing_pk_cols}") 

76 

77 # Strategy 2: Combine all unique constraints (union approach) 

78 all_unique_columns = [] 

79 for constraint in table_schema.unique_constraints: 

80 # Only include constraint columns that are available in matched_columns 

81 available_constraint_cols = [col for col in constraint.columns if col in available_columns] 

82 all_unique_columns.extend(available_constraint_cols) 

83 

84 # Remove duplicates 

85 unique_columns = list(set(all_unique_columns)) 

86 

87 if unique_columns: 87 ↛ 97line 87 didn't jump to line 97 because the condition on line 87 was always true

88 strategy = ConflictStrategy( 

89 type="UNIQUE_COMBINED", 

90 columns=unique_columns, 

91 description=f"Combined unique constraints conflict resolution on: {unique_columns}" 

92 ) 

93 logger.debug(f"Using UNIQUE_COMBINED strategy: {strategy.description}") 

94 return strategy 

95 

96 # Strategy 3: Fallback to INSERT only 

97 strategy = ConflictStrategy( 

98 type="INSERT_ONLY", 

99 columns=[], 

100 description="No conflict resolution - INSERT only (no unique constraints found)" 

101 ) 

102 logger.warning(f"Fallback to INSERT_ONLY strategy: {strategy.description}") 

103 return strategy 

104 

105 

106def deduplicate_temp_table( 

107 connection, 

108 temp_table_name: str, 

109 conflict_columns: list[str] 

110) -> DeduplicationResult: 

111 """ 

112 Deduplicate temp table based on conflict columns, keeping last occurrence. 

113 Creates a new cleaned temp table and logs dropped rows. 

114 

115 Args: 

116 connection: Database connection 

117 temp_table_name: Name of the source temp table 

118 conflict_columns: Columns to use for deduplication 

119 schema_name: Schema name (default: 'public') 

120 

121 Returns: 

122 DeduplicationResult with statistics 

123 

124 Raises: 

125 PgsqlUpserterError: If deduplication fails 

126 """ 

127 logger.debug(f"Starting deduplication of '{temp_table_name}' on columns: {conflict_columns}") 

128 

129 try: 

130 with connection.cursor() as cursor: 

131 # Get original count 

132 cursor.execute(f"SELECT COUNT(*) FROM {temp_table_name}") 

133 original_count = cursor.fetchone()[0] 

134 logger.debug(f"Original temp table count: {original_count} rows") 

135 

136 if not conflict_columns: 136 ↛ 138line 136 didn't jump to line 138 because the condition on line 136 was never true

137 # No deduplication needed for INSERT_ONLY strategy 

138 return DeduplicationResult( 

139 original_count=original_count, 

140 deduplicated_count=original_count, 

141 dropped_count=0, 

142 dropped_reasons={} 

143 ) 

144 

145 # Create cleaned temp table name 

146 cleaned_table_name = f"{temp_table_name}_cleaned" 

147 

148 # Step 1: Build NULL checking conditions for conflict columns 

149 # We need to handle text columns differently from other types 

150 cursor.execute(f""" 

151 SELECT column_name, data_type 

152 FROM information_schema.columns 

153 WHERE table_name = '{temp_table_name.split('.')[-1]}' 

154 AND column_name = ANY(%s) 

155 """, (conflict_columns,)) 

156 

157 column_types = {row[0]: row[1] for row in cursor.fetchall()} 

158 

159 null_conditions = [] 

160 for col in conflict_columns: 

161 if col in column_types: 161 ↛ 171line 161 didn't jump to line 171 because the condition on line 161 was always true

162 data_type = column_types[col] 

163 if data_type in ('text', 'varchar', 'character varying', 'char'): 

164 # For text columns, check both NULL and empty string 

165 null_conditions.append(f"({col} IS NULL OR {col} = '')") 

166 else: 

167 # For other types (date, numeric, etc.), only check NULL 

168 null_conditions.append(f"{col} IS NULL") 

169 else: 

170 # Fallback: assume text type 

171 null_conditions.append(f"({col} IS NULL OR {col} = '')") 

172 

173 null_where_clause = " OR ".join(null_conditions) 

174 

175 # Count rows with NULLs before removing 

176 cursor.execute(f""" 

177 SELECT COUNT(*) 

178 FROM {temp_table_name} 

179 WHERE {null_where_clause} 

180 """) 

181 null_count = cursor.fetchone()[0] 

182 

183 # Step 2: Create cleaned table with deduplication (keeping last occurrence) 

184 conflict_columns_str = ", ".join(conflict_columns) 

185 

186 cursor.execute(f""" 

187 CREATE TEMP TABLE {cleaned_table_name} AS 

188 SELECT * FROM ( 

189 SELECT *, 

190 ROW_NUMBER() OVER ( 

191 PARTITION BY {conflict_columns_str} 

192 ORDER BY ctid DESC 

193 ) as rn 

194 FROM {temp_table_name} 

195 WHERE NOT ({null_where_clause}) 

196 ) ranked 

197 WHERE rn = 1 

198 """) 

199 

200 # Get final count 

201 cursor.execute(f"SELECT COUNT(*) FROM {cleaned_table_name}") 

202 deduplicated_count = cursor.fetchone()[0] 

203 

204 # Calculate dropped counts 

205 duplicate_count = (original_count - null_count) - deduplicated_count 

206 total_dropped = original_count - deduplicated_count 

207 

208 dropped_reasons = {} 

209 if null_count > 0: 209 ↛ 210line 209 didn't jump to line 210 because the condition on line 209 was never true

210 dropped_reasons["null_or_empty_conflict_columns"] = null_count 

211 if duplicate_count > 0: 

212 dropped_reasons["duplicate_conflict_keys"] = duplicate_count 

213 

214 # Drop original temp table and rename cleaned table 

215 cursor.execute(f"DROP TABLE {temp_table_name}") 

216 cursor.execute(f"ALTER TABLE {cleaned_table_name} DROP COLUMN rn") 

217 cursor.execute(f"ALTER TABLE {cleaned_table_name} RENAME TO {temp_table_name.split('.')[-1]}") 

218 

219 result = DeduplicationResult( 

220 original_count=original_count, 

221 deduplicated_count=deduplicated_count, 

222 dropped_count=total_dropped, 

223 dropped_reasons=dropped_reasons 

224 ) 

225 

226 logger.debug(f"Deduplication completed: {original_count} → {deduplicated_count} rows " 

227 f"({total_dropped} dropped)") 

228 for reason, count in dropped_reasons.items(): 

229 logger.debug(f" - {reason}: {count} rows") 

230 

231 return result 

232 

233 except Exception as e: 

234 logger.error(f"Deduplication failed: {e}") 

235 raise PgsqlUpserterError(f"Failed to deduplicate temp table: {e}") from e 

236 

237 

238def execute_upsert( 

239 connection, 

240 temp_table_name: str, 

241 target_table: str, 

242 conflict_strategy: ConflictStrategy, 

243 matched_columns: list[str], 

244 schema_name: str = 'public' 

245) -> tuple[int, int]: 

246 """ 

247 Execute the final upsert operation using INSERT...ON CONFLICT. 

248 

249 Args: 

250 connection: Database connection 

251 temp_table_name: Name of the cleaned temp table 

252 target_table: Target table name 

253 conflict_strategy: Conflict resolution strategy 

254 matched_columns: List of columns available in the temp table 

255 schema_name: Schema name (default: 'public') 

256 

257 Returns: 

258 Tuple of (rows_inserted, rows_updated) 

259 

260 Raises: 

261 PgsqlUpserterError: If upsert operation fails 

262 """ 

263 logger.debug(f"Executing upsert from '{temp_table_name}' to '{schema_name}.{target_table}' " 

264 f"using {conflict_strategy.type} strategy") 

265 

266 try: 

267 with connection.cursor() as cursor: 

268 # Use matched_columns instead of querying target table columns 

269 columns_str = ", ".join(matched_columns) 

270 

271 if conflict_strategy.type == "INSERT_ONLY": 271 ↛ 273line 271 didn't jump to line 273 because the condition on line 271 was never true

272 # Simple INSERT without conflict resolution 

273 cursor.execute(f""" 

274 INSERT INTO {schema_name}.{target_table} ({columns_str}) 

275 SELECT {columns_str} 

276 FROM {temp_table_name} 

277 """) 

278 

279 rows_affected = cursor.rowcount 

280 logger.debug(f"INSERT_ONLY completed: {rows_affected} rows inserted") 

281 return rows_affected, 0 

282 

283 else: 

284 # INSERT...ON CONFLICT with UPDATE 

285 conflict_columns_str = ", ".join(conflict_strategy.columns) 

286 

287 # Build UPDATE SET clause (matched columns except conflict columns) 

288 update_columns = [col for col in matched_columns if col not in conflict_strategy.columns] 

289 update_set_clause = ", ".join([f"{col} = EXCLUDED.{col}" for col in update_columns]) 

290 

291 # Track before counts for calculating inserts vs updates 

292 cursor.execute(f"SELECT COUNT(*) FROM {schema_name}.{target_table}") 

293 before_count = cursor.fetchone()[0] 

294 

295 cursor.execute(f""" 

296 INSERT INTO {schema_name}.{target_table} ({columns_str}) 

297 SELECT {columns_str} 

298 FROM {temp_table_name} 

299 ON CONFLICT ({conflict_columns_str}) 

300 DO UPDATE SET {update_set_clause} 

301 """) 

302 

303 # Calculate insert vs update counts 

304 cursor.execute(f"SELECT COUNT(*) FROM {schema_name}.{target_table}") 

305 after_count = cursor.fetchone()[0] 

306 

307 cursor.execute(f"SELECT COUNT(*) FROM {temp_table_name}") 

308 temp_count = cursor.fetchone()[0] 

309 

310 rows_inserted = after_count - before_count 

311 rows_updated = temp_count - rows_inserted 

312 

313 logger.debug(f"Upsert completed: {rows_inserted} inserted, {rows_updated} updated") 

314 return rows_inserted, rows_updated 

315 

316 except Exception as e: 

317 logger.error(f"Upsert operation failed: {e}") 

318 raise PgsqlUpserterError(f"Failed to execute upsert: {e}") from e