Coverage for pgsql_upserter / upsert_engine.py: 94.05%
78 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-26 18:15 -0300
1"""Main upsert engine - public API for PostgreSQL upsert operations."""
3import csv
4import logging
5import psycopg2
7from dataclasses import dataclass
8from pathlib import Path
10from .schema_inspector import inspect_table_schema
11from .column_matcher import match_columns
12from .temp_staging import create_temp_table, bulk_insert_to_temp
13from .conflict_resolver import (
14 find_conflict_strategy,
15 deduplicate_temp_table,
16 execute_upsert,
17 DeduplicationResult,
18 ConflictStrategy
19)
20from .config import create_connection_from_env, test_connection
22# Configure module logger
23logger = logging.getLogger(__name__)
26@dataclass
27class UpsertResult:
28 """Complete results from upsert workflow execution."""
29 rows_inserted: int
30 rows_updated: int
31 total_affected: int
32 deduplication_result: DeduplicationResult
33 matched_columns: list[str]
34 conflict_strategy_type: str
35 conflict_strategy_description: str
38@staticmethod
39def read_csv_to_dict_list(csv_path: str | Path) -> list[dict[str, str]]:
40 """Read CSV file and return list of dictionaries.
42 Args:
43 csv_path: Path to the CSV file
45 Returns:
46 List of dictionaries with CSV data
48 Note:
49 CSV data is returned as strings. Type conversion should be handled
50 by the database driver during insertion.
51 """
52 csv_file = Path(csv_path)
53 data_list = []
54 with open(csv_file, 'r', encoding='utf-8') as f:
55 reader = csv.DictReader(f)
56 data_list = list(reader)
57 return data_list
60@staticmethod
61def execute_upsert_workflow(
62 connection: psycopg2.extensions.connection,
63 data: list[dict[str, str]] | str | Path,
64 target_table: str,
65 conflict_columns: list[str] | None = None,
66 update_columns: list[str] | None = None,
67 batch_size: int = 1000,
68 keep_temp_table: bool = False,
69 schema: str = 'public'
70) -> UpsertResult:
71 """Execute complete upsert workflow with automatic conflict detection.
73 This function orchestrates the complete upsert process including:
74 1. Data input handling (CSV files or direct data lists)
75 2. Temporary table creation with automatic column detection
76 3. Conflict strategy detection (primary key, unique constraints, or insert-only)
77 4. Data deduplication using appropriate strategy
78 5. Upsert execution with proper conflict resolution
80 Args:
81 connection: Active PostgreSQL database connection
82 target_table: Name of the target table for upsert operation
83 data: Input data as list of dictionaries, CSV file path, or Path object
84 conflict_columns: Optional override for conflict detection columns. If provided,
85 these columns will be used for conflict resolution instead of
86 automatic detection (primary keys, unique constraints)
87 update_columns: Optional override for columns to update on conflict. If not
88 provided, all matched columns will be updated
89 batch_size: Number of rows to process in each batch during temp table population
90 temp_table_prefix: Prefix for temporary table name (default: "_temp_")
91 keep_temp_table: Whether to preserve temporary table after operation
93 Returns:
94 UpsertResult: Object containing operation results and statistics
96 Raises:
97 ValueError: If data is empty or target_table is invalid
98 psycopg2.Error: For database connection or operation errors
100 Example:
101 >>> # Using direct data
102 >>> data = [{'id': 1, 'name': 'John'}, {'id': 2, 'name': 'Jane'}]
103 >>> result = execute_upsert_workflow(conn, 'users', data)
104 >>> print(f"Inserted: {result.inserted_count}, Updated: {result.updated_count}")
106 >>> # Using CSV file
107 >>> result = execute_upsert_workflow(conn, 'users', 'data.csv')
109 Note:
110 For serverless ETL: This function is designed to work seamlessly with
111 API responses by accepting direct data lists, eliminating the need
112 for intermediate CSV file creation in lambda/cloud functions.
113 """
114 logger.info(f"Starting upsert workflow for table '{target_table}'")
116 # Step 1: Handle input data
117 if isinstance(data, (str, Path)): 117 ↛ 118line 117 didn't jump to line 118 because the condition on line 117 was never true
118 logger.info(f"Reading CSV file: {data}")
119 data_list = read_csv_to_dict_list(data)
120 else:
121 data_list = data
123 if not data_list:
124 raise ValueError("No data provided for upsert operation")
126 logger.info(f"Processing {len(data_list)} rows")
128 # Step 2: Inspect target table schema
129 target_schema = inspect_table_schema(connection, target_table, schema)
130 logger.info("Target table schema inspected")
132 # Step 3: Match and map columns
133 column_mapping = match_columns(data_list, target_schema)
134 matched_columns = column_mapping['matched_columns']
135 logger.info(f"Matched columns: {matched_columns}")
137 # Step 4: Create and populate temp table
138 # Create temp table with auto-generated name
139 temp_table_name = create_temp_table(connection, target_table, schema)
140 logger.info(f"Created temp table: {temp_table_name}")
142 # Populate temp table
143 rows_inserted = bulk_insert_to_temp(
144 connection=connection,
145 temp_table_name=temp_table_name,
146 data_list=data_list,
147 matched_columns=matched_columns,
148 target_schema=target_schema,
149 batch_size=batch_size
150 )
151 logger.info(f"Populated temp table with {rows_inserted} rows")
153 try:
154 # Step 5: Find conflict strategy
155 if conflict_columns:
156 # Use user-provided conflict columns
157 conflict_strategy = ConflictStrategy(
158 type="USER_DEFINED",
159 columns=conflict_columns,
160 description=f"User-defined conflict resolution on: {conflict_columns}"
161 )
162 logger.info(f"Using user-defined conflict strategy: {conflict_strategy.description}")
163 else:
164 # Automatic detection
165 conflict_strategy = find_conflict_strategy(
166 target_schema,
167 matched_columns
168 )
169 logger.info(f"Using automatic conflict strategy: {conflict_strategy.type}")
171 # Step 6: Deduplicate temp table
172 dedup_result = deduplicate_temp_table(
173 connection,
174 temp_table_name,
175 conflict_strategy.columns
176 )
177 logger.info(f"Deduplication: {dedup_result.original_count} -> {dedup_result.deduplicated_count}")
179 # Step 7: Execute upsert
180 columns_to_update = update_columns or matched_columns or list(data_list[0].keys())
181 inserted_count, updated_count = execute_upsert(
182 connection,
183 temp_table_name,
184 target_table,
185 conflict_strategy,
186 columns_to_update,
187 schema,
188 )
189 logger.info(f"Upsert complete: {inserted_count} inserted, {updated_count} updated\n")
191 # Step 8: Create final result
192 result = UpsertResult(
193 rows_inserted=inserted_count,
194 rows_updated=updated_count,
195 total_affected=inserted_count + updated_count,
196 deduplication_result=dedup_result,
197 matched_columns=matched_columns or list(data_list[0].keys()),
198 conflict_strategy_type=conflict_strategy.type,
199 conflict_strategy_description=conflict_strategy.description
200 )
202 return result
204 finally:
205 # Clean up temp table unless requested to keep
206 if not keep_temp_table:
207 try:
208 with connection.cursor() as cursor:
209 cursor.execute(f"DROP TABLE IF EXISTS {temp_table_name}")
210 connection.commit()
211 logger.debug(f"Cleaned up temp table: {temp_table_name}")
212 except Exception as e:
213 logger.warning(f"Failed to clean up temp table {temp_table_name}: {e}")
216class UpsertEngine:
217 """Main interface for PostgreSQL upsert operations.
219 This class provides a clean API for all upsert functionality including
220 connection management, data loading, and workflow execution.
221 """
223 @staticmethod
224 def create_connection() -> psycopg2.extensions.connection:
225 """Create PostgreSQL connection from environment variables.
227 Expected environment variables:
228 - PGHOST (default: localhost)
229 - PGPORT (default: 5432)
230 - PGDATABASE (required)
231 - PGUSER (required)
232 - PGPASSWORD (required)
234 Returns:
235 psycopg2.connection: Active database connection
237 Raises:
238 ConnectionError: If connection fails
239 """
240 return create_connection_from_env()
242 @staticmethod
243 def test_connection() -> psycopg2.extensions.connection:
244 """Test database connection and validate permissions for upsert operations.
246 This method creates a connection and validates that the user has
247 sufficient permissions to perform upsert operations (CREATE TEMP TABLE).
249 Returns:
250 psycopg2.connection: Validated connection ready for use
252 Raises:
253 ConnectionError: If connection fails
254 PermissionError: If user lacks required permissions
255 """
256 return test_connection()
258 @staticmethod
259 def upsert_data(
260 connection: psycopg2.extensions.connection,
261 data: list[dict[str, str]] | str | Path,
262 target_table: str,
263 conflict_columns: list[str] | None = None,
264 update_columns: list[str] | None = None,
265 batch_size: int = 1000,
266 keep_temp_table: bool = False,
267 schema: str = 'public'
268 ) -> UpsertResult:
269 """Execute complete upsert workflow with automatic conflict detection.
271 This function orchestrates the complete upsert process including:
272 1. Data input handling (CSV files or direct data lists)
273 2. Temporary table creation with automatic column detection
274 3. Conflict strategy detection (primary key, unique constraints, or insert-only)
275 4. Data deduplication using appropriate strategy
276 5. Upsert execution with proper conflict resolution
278 Args:
279 connection: Active PostgreSQL database connection
280 target_table: Name of the target table for upsert operation
281 data: Input data as list of dictionaries, CSV file path, or Path object
282 conflict_columns: Optional override for conflict detection columns. If provided,
283 these columns will be used for conflict resolution instead of
284 automatic detection (primary keys, unique constraints)
285 update_columns: Optional override for columns to update on conflict. If not
286 provided, all matched columns will be updated
287 batch_size: Number of rows to process in each batch during temp table population
288 temp_table_prefix: Prefix for temporary table name (default: "_temp_")
289 keep_temp_table: Whether to preserve temporary table after operation
291 Returns:
292 UpsertResult: Object containing operation results and statistics
294 Raises:
295 ValueError: If data is empty or target_table is invalid
296 psycopg2.Error: For database connection or operation errors
298 Example:
299 >>> # Using direct data
300 >>> data = [{'id': 1, 'name': 'John'}, {'id': 2, 'name': 'Jane'}]
301 >>> result = execute_upsert_workflow(conn, 'users', data)
302 >>> print(f"Inserted: {result.inserted_count}, Updated: {result.updated_count}")
304 >>> # Using CSV file
305 >>> result = execute_upsert_workflow(conn, 'users', 'data.csv')
307 Note:
308 For serverless ETL: This function is designed to work seamlessly with
309 API responses by accepting direct data lists, eliminating the need
310 for intermediate CSV file creation in lambda/cloud functions.
311 """
313 return execute_upsert_workflow(
314 connection=connection,
315 data=data,
316 target_table=target_table,
317 conflict_columns=conflict_columns,
318 update_columns=update_columns,
319 batch_size=batch_size,
320 keep_temp_table=keep_temp_table,
321 schema=schema
322 )