sqlalchemy-upsert-kit 0.1.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,156 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ import typing as T
4
+ import dataclasses
5
+
6
+ import sqlalchemy as sa
7
+
8
+ from ..exc import UpsertTestError
9
+
10
+ from .executor import UpsertExecutor
11
+
12
+
13
+ @dataclasses.dataclass
14
+ class InsertOrIgnoreExecutor(UpsertExecutor):
15
+ def apply_strategy(
16
+ self,
17
+ conn: sa.Connection,
18
+ trans: sa.Transaction,
19
+ ):
20
+ if self._raise_on_target_insert: # Testing flag
21
+ raise UpsertTestError("error on target insertion")
22
+ stmt = self.table.insert().from_select(
23
+ list(self._temp_table.columns.keys()),
24
+ sa.select(self._temp_table)
25
+ .select_from(
26
+ self._temp_table.outerjoin( # LEFT JOIN to find non-matches
27
+ self.table,
28
+ self._temp_table.c[self.pk_name] == self.table.c[self.pk_name],
29
+ )
30
+ )
31
+ .where(
32
+ self.table.c[self.pk_name].is_(None)
33
+ ), # Only insert where no match exists
34
+ )
35
+ res = conn.execute(stmt)
36
+ try:
37
+ self._inserted_rows = res.rowcount if res.rowcount is not None else 0
38
+ self._ignored_rows = len(self.values) - self._inserted_rows
39
+ except: # pragma: no cover
40
+ self._inserted_rows = 0
41
+ self._ignored_rows = len(self.values)
42
+
43
+
44
+ def insert_or_ignore(
45
+ engine: sa.Engine,
46
+ table: sa.Table,
47
+ values: list[dict[str, T.Any]],
48
+ metadata: T.Optional[sa.MetaData] = None,
49
+ temp_table_name: T.Optional[str] = None,
50
+ conn: T.Optional[sa.Connection] = None,
51
+ trans: T.Optional[sa.Transaction] = None,
52
+ _raise_on_temp_table_create: bool = False,
53
+ _raise_on_temp_data_insert: bool = False,
54
+ _raise_on_target_insert: bool = False,
55
+ _raise_on_temp_table_drop: bool = False,
56
+ ) -> tuple[int, int]:
57
+ """
58
+ Perform high-performance bulk INSERT-IF-NOT-EXISTS operation using temporary table.
59
+
60
+ This function performs conditional bulk insertion: only inserts records whose
61
+ primary keys don't already exist in the target table. This is equivalent to
62
+ "INSERT IGNORE" or "INSERT ... ON CONFLICT DO NOTHING" but works more
63
+ efficiently.
64
+
65
+ **Algorithm**:
66
+
67
+ 1. Creates temporary table and loads all candidate data
68
+ 2. Uses LEFT JOIN to identify records not in target table
69
+ 3. Bulk inserts only the non-conflicting records
70
+ 4. Cleans up temporary resources
71
+
72
+ This approach is ideal for:
73
+
74
+ - Incremental data loading where duplicates should be ignored
75
+ - ETL processes that need idempotent behavior
76
+ - Syncing data from external sources
77
+
78
+ **Transaction Management**:
79
+
80
+ This function supports both auto-managed and user-managed transaction modes.
81
+ See the module-level documentation for detailed explanations of each mode.
82
+
83
+ :param engine: SQLAlchemy engine for database connection
84
+ :param table: Target table for conditional insertion
85
+ :param values: Records to insert if they don't exist.
86
+ Must include primary key values for conflict detection.
87
+ :param metadata: Optional metadata instance for temporary table isolation.
88
+ If None, a new MetaData instance is created for clean separation.
89
+ :param temp_table_name: Optional custom name for temporary table.
90
+ If None, generates unique name with timestamp to avoid conflicts.
91
+ :param conn: Optional database connection for user-managed transaction mode.
92
+ Must be provided together with ``trans`` parameter.
93
+ :param trans: Optional transaction for user-managed transaction mode.
94
+ Must be provided together with ``conn`` parameter.
95
+
96
+ :returns: Tuple of (ignored_rows, inserted_rows):
97
+ - ignored_rows: Number of records that were not inserted (already existed)
98
+ - inserted_rows: Number of new records successfully inserted
99
+
100
+ :raises ValueError: When conn and trans parameters are provided inconsistently
101
+ (one is None while the other is not)
102
+ :raises UpsertTestError: When testing flags are enabled and corresponding operations fail
103
+
104
+ **Examples**:
105
+
106
+ Auto-managed transaction (default mode)::
107
+
108
+ # Function manages its own transaction
109
+ ignored, inserted = insert_or_ignore(engine, users_table, new_data)
110
+
111
+ User-managed transaction mode::
112
+
113
+ # Operation is part of larger transaction
114
+ with engine.connect() as conn:
115
+ with conn.begin() as trans:
116
+ # Other operations...
117
+ ignored, inserted = insert_or_ignore(
118
+ engine, users_table, new_data, conn=conn, trans=trans
119
+ )
120
+ # More operations...
121
+
122
+ Conflict detection example::
123
+
124
+ # Target table has records with id=1,2,3
125
+ new_data = [
126
+ {'id': 2, 'name': 'Bob'}, # Exists - will be ignored
127
+ {'id': 4, 'name': 'Charlie'}, # New - will be inserted
128
+ {'id': 5, 'name': 'David'}, # New - will be inserted
129
+ ]
130
+ ignored, inserted = insert_or_ignore(engine, users_table, new_data)
131
+ # Result: ignored=1, inserted=2
132
+
133
+ .. note::
134
+
135
+ Parameters prefixed with ``_raise_on_`` are exclusively for testing error
136
+ handling and cleanup behavior. Never use these in production code.
137
+ """
138
+ if not values: # pragma: no cover
139
+ return 0, 0 # No-op for empty data
140
+
141
+ executor = InsertOrIgnoreExecutor.new(
142
+ engine=engine,
143
+ table=table,
144
+ values=values,
145
+ metadata=metadata,
146
+ temp_table_name=temp_table_name,
147
+ conn=conn,
148
+ trans=trans,
149
+ _raise_on_temp_table_create=_raise_on_temp_table_create,
150
+ _raise_on_temp_data_insert=_raise_on_temp_data_insert,
151
+ _raise_on_target_insert=_raise_on_target_insert,
152
+ _raise_on_temp_table_drop=_raise_on_temp_table_drop,
153
+ )
154
+ executor.run()
155
+ ignored_rows, inserted_rows = executor._ignored_rows, executor._inserted_rows
156
+ return ignored_rows, inserted_rows
@@ -0,0 +1,254 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ import typing as T
4
+ import dataclasses
5
+
6
+ import sqlalchemy as sa
7
+
8
+ from ..exc import UpsertTestError
9
+
10
+ from .executor import UpsertExecutor
11
+
12
+
13
+ @dataclasses.dataclass
14
+ class InsertOrMergeExecutor(UpsertExecutor):
15
+ def apply_strategy(
16
+ self,
17
+ conn: sa.Connection,
18
+ trans: sa.Transaction,
19
+ ):
20
+ """
21
+ Apply merge strategy: update existing records with selected columns, insert new records.
22
+
23
+ This implementation follows the merge algorithm:
24
+ 1. Create temp table t3 with merged data (existing records with selective updates)
25
+ 2. Delete conflicting records from target table t1
26
+ 3. Insert both updated records (from t3) and new records (from t2) into t1
27
+ """
28
+ if self._raise_on_target_delete: # Testing flag
29
+ raise UpsertTestError("error on target deletion")
30
+
31
+ # Step 1: Create merge temp table (t3) with updated data for existing records
32
+ merge_temp_table_name = f"{self.temp_table_name}_merge"
33
+ merge_temp_table = sa.Table(
34
+ merge_temp_table_name,
35
+ self.metadata,
36
+ *[
37
+ sa.Column(col.name, col.type, nullable=col.nullable)
38
+ for col in self.table.columns
39
+ ],
40
+ prefixes=["TEMPORARY"],
41
+ )
42
+ merge_temp_table.create(conn)
43
+
44
+ try:
45
+ if self._raise_on_merge_update: # Testing flag
46
+ raise UpsertTestError("error on merge update")
47
+
48
+ # Build the merge query: existing records with selective column updates
49
+ # For columns in self.columns: use values from temp table (t2)
50
+ # For columns not in self.columns: use values from target table (t1)
51
+ select_cols = []
52
+ for col in self.table.columns:
53
+ if col.name in self.columns:
54
+ # Use new value from temp table
55
+ select_cols.append(self._temp_table.c[col.name].label(col.name))
56
+ else:
57
+ # Preserve existing value from target table
58
+ select_cols.append(self.table.c[col.name].label(col.name))
59
+
60
+ # Create merged records for existing/conflicting records
61
+ merge_stmt = merge_temp_table.insert().from_select(
62
+ [col.name for col in self.table.columns],
63
+ sa.select(*select_cols).select_from(
64
+ self.table.join(
65
+ self._temp_table,
66
+ self.table.c[self.pk_name] == self._temp_table.c[self.pk_name],
67
+ )
68
+ ),
69
+ )
70
+ merge_result = conn.execute(merge_stmt)
71
+ self._updated_rows = (
72
+ merge_result.rowcount if merge_result.rowcount is not None else 0
73
+ )
74
+
75
+ # Step 2: Delete existing records that will be replaced with merged data
76
+ inner = sa.select(self.table.c[self.pk_name]).join(
77
+ self._temp_table,
78
+ self.table.c[self.pk_name] == self._temp_table.c[self.pk_name],
79
+ )
80
+ delete_stmt = self.table.delete().where(
81
+ self.table.c[self.pk_name].in_(inner)
82
+ )
83
+ conn.execute(delete_stmt)
84
+
85
+ # Step 3: Insert merged records (updated existing records)
86
+ if self._raise_on_target_insert: # Testing flag
87
+ raise UpsertTestError("error on target insertion")
88
+
89
+ if self._updated_rows > 0:
90
+ insert_merged_stmt = self.table.insert().from_select(
91
+ list(merge_temp_table.columns.keys()),
92
+ sa.select(*list(merge_temp_table.columns.values())),
93
+ )
94
+ conn.execute(insert_merged_stmt)
95
+
96
+ # Step 4: Insert new records (records that don't exist in target table)
97
+ # Use LEFT JOIN to find records in temp table that don't exist in target
98
+ new_records_stmt = self.table.insert().from_select(
99
+ list(self._temp_table.columns.keys()),
100
+ sa.select(*list(self._temp_table.columns.values()))
101
+ .select_from(
102
+ self._temp_table.outerjoin(
103
+ self.table,
104
+ self._temp_table.c[self.pk_name] == self.table.c[self.pk_name],
105
+ )
106
+ )
107
+ .where(self.table.c[self.pk_name].is_(None)),
108
+ )
109
+ insert_result = conn.execute(new_records_stmt)
110
+ self._inserted_rows = (
111
+ insert_result.rowcount if insert_result.rowcount is not None else 0
112
+ )
113
+
114
+ finally:
115
+ # Clean up merge temp table
116
+ merge_temp_table.drop(conn)
117
+ self.metadata.remove(merge_temp_table)
118
+
119
+
120
+ def insert_or_merge(
121
+ engine: sa.Engine,
122
+ table: sa.Table,
123
+ values: list[dict[str, T.Any]],
124
+ columns: list[str],
125
+ metadata: T.Optional[sa.MetaData] = None,
126
+ temp_table_name: T.Optional[str] = None,
127
+ conn: T.Optional[sa.Connection] = None,
128
+ trans: T.Optional[sa.Transaction] = None,
129
+ _raise_on_temp_table_create: bool = False,
130
+ _raise_on_temp_data_insert: bool = False,
131
+ _raise_on_target_delete: bool = False,
132
+ _raise_on_target_insert: bool = False,
133
+ _raise_on_temp_table_drop: bool = False,
134
+ _raise_on_merge_update: bool = False,
135
+ ) -> tuple[int, int]:
136
+ """
137
+ Perform high-performance bulk INSERT-OR-MERGE operation using temporary table.
138
+
139
+ This function performs bulk merge operations: updates specific columns of existing
140
+ records while preserving other columns, and inserts records that don't exist.
141
+ This is ideal for incremental data updates where only certain fields need updating.
142
+
143
+ **Algorithm**:
144
+
145
+ 1. Creates temporary table (t2) and loads all candidate data
146
+ 2. Creates merge temporary table (t3) with selective column updates
147
+ 3. Generates merged records by combining existing data with new values for specified columns
148
+ 4. Deletes conflicting records from target table (t1)
149
+ 5. Inserts merged records and new records into target table
150
+ 6. Cleans up temporary resources
151
+
152
+ This approach is ideal for:
153
+
154
+ - Selective column updates (e.g., updating timestamps while preserving descriptions)
155
+ - Incremental data synchronization
156
+ - Preserving historical data in non-updated columns
157
+
158
+ **Transaction Management**:
159
+
160
+ This function supports both auto-managed and user-managed transaction modes.
161
+ See the module-level documentation for detailed explanations of each mode.
162
+
163
+ :param engine: SQLAlchemy engine for database connection
164
+ :param table: Target table for upsert operation
165
+ :param values: Records to merge.
166
+ Must include primary key values for conflict detection.
167
+ :param columns: List of columns to update with new values. Other columns will
168
+ remain unchanged for existing records. This is required, otherwise user should use
169
+ insert_or_ignore or insert_or_replace strategy.
170
+ :param metadata: Optional metadata instance for temporary table isolation.
171
+ If None, a new MetaData instance is created for clean separation.
172
+ :param temp_table_name: Optional custom name for temporary table.
173
+ If None, generates unique name with timestamp to avoid conflicts.
174
+ :param conn: Optional database connection for user-managed transaction mode.
175
+ Must be provided together with ``trans`` parameter.
176
+ :param trans: Optional transaction for user-managed transaction mode.
177
+ Must be provided together with ``conn`` parameter.
178
+
179
+ :returns: Tuple of (updated_rows, inserted_rows):
180
+ - updated_rows: Number of existing records that were updated with new column values
181
+ - inserted_rows: Number of new records that were inserted
182
+
183
+ :raises ValueError: When conn and trans parameters are provided inconsistently
184
+ or when columns parameter is empty
185
+ :raises UpsertTestError: When testing flags are enabled and corresponding operations fail
186
+
187
+ **Examples**:
188
+
189
+ Auto-managed transaction (default mode)::
190
+
191
+ # Function manages its own transaction
192
+ updated, inserted = merge(engine, users_table, new_data, columns=["update_at"])
193
+
194
+ User-managed transaction mode::
195
+
196
+ # Operation is part of larger transaction
197
+ with engine.connect() as conn:
198
+ with conn.begin() as trans:
199
+ # Other operations...
200
+ updated, inserted = merge(
201
+ engine, users_table, new_data, columns=["update_at"],
202
+ conn=conn, trans=trans
203
+ )
204
+ # More operations...
205
+
206
+ Selective column update example::
207
+
208
+ # Target table has records with id=1,2,3
209
+ new_data = [
210
+ {'id': 2, 'name': 'Bob', 'update_at': '2024-01-02'}, # Exists - will be merged
211
+ {'id': 4, 'name': 'Charlie', 'update_at': '2024-01-02'}, # New - will be inserted
212
+ ]
213
+ updated, inserted = merge(engine, users_table, new_data, columns=["update_at"])
214
+ # Result: updated=1, inserted=1
215
+ # Record id=2 will have new update_at but original name preserved
216
+
217
+ .. note::
218
+
219
+ Parameters prefixed with ``_raise_on_`` are exclusively for testing error
220
+ handling and cleanup behavior. Never use these in production code.
221
+
222
+ .. warning::
223
+
224
+ The columns parameter must not be empty. If you want to update all columns,
225
+ use insert_or_replace instead. If you want to ignore conflicts, use
226
+ insert_or_ignore instead.
227
+ """
228
+ if not values: # pragma: no cover
229
+ return 0, 0 # No-op for empty data
230
+
231
+ if not columns: # pragma: no cover
232
+ raise ValueError(
233
+ "columns parameter cannot be empty. Use insert_or_replace for full updates or insert_or_ignore for conflict-free inserts."
234
+ )
235
+
236
+ executor = InsertOrMergeExecutor.new(
237
+ engine=engine,
238
+ table=table,
239
+ values=values,
240
+ metadata=metadata,
241
+ temp_table_name=temp_table_name,
242
+ conn=conn,
243
+ trans=trans,
244
+ columns=columns,
245
+ _raise_on_temp_table_create=_raise_on_temp_table_create,
246
+ _raise_on_temp_data_insert=_raise_on_temp_data_insert,
247
+ _raise_on_target_delete=_raise_on_target_delete,
248
+ _raise_on_target_insert=_raise_on_target_insert,
249
+ _raise_on_temp_table_drop=_raise_on_temp_table_drop,
250
+ _raise_on_merge_update=_raise_on_merge_update,
251
+ )
252
+ executor.run()
253
+ updated_rows, inserted_rows = executor._updated_rows, executor._inserted_rows
254
+ return updated_rows, inserted_rows
@@ -0,0 +1,177 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ import typing as T
4
+ import dataclasses
5
+
6
+ import sqlalchemy as sa
7
+
8
+ from ..exc import UpsertTestError
9
+
10
+ from .executor import UpsertExecutor
11
+
12
+
13
+ @dataclasses.dataclass
14
+ class InsertOrReplaceExcutor(UpsertExecutor):
15
+ def apply_strategy(
16
+ self,
17
+ conn: sa.Connection,
18
+ trans: sa.Transaction,
19
+ ):
20
+ # Step 3: Delete existing records that will be replaced
21
+ # Uses JOIN to identify conflicting records for optimal performance
22
+ if self._raise_on_target_delete: # Testing flag
23
+ raise UpsertTestError("error on target deletion")
24
+ inner = sa.select(self.table.c[self.pk_name]).join(
25
+ self._temp_table,
26
+ self.table.c[self.pk_name] == self._temp_table.c[self.pk_name],
27
+ )
28
+ stmt = self.table.delete().where(self.table.c[self.pk_name].in_(inner))
29
+ res = conn.execute(stmt)
30
+ try:
31
+ self._replaced_rows = res.rowcount if res.rowcount is not None else 0
32
+ except: # pragma: no cover
33
+ self._replaced_rows = 0
34
+
35
+ # Step 4: Insert all records from temp table (both replacements and new)
36
+ # This includes records that replace deleted ones and completely new records
37
+ if self._raise_on_target_insert: # Testing flag
38
+ raise UpsertTestError("error on target insertion")
39
+ stmt = self.table.insert().from_select(
40
+ list(self._temp_table.columns.keys()),
41
+ sa.select(*list(self._temp_table.columns.values())),
42
+ )
43
+ res = conn.execute(stmt)
44
+ try:
45
+ total_inserted = res.rowcount if res.rowcount is not None else 0
46
+ self._inserted_rows = total_inserted - self._replaced_rows
47
+ except: # pragma: no cover
48
+ self._inserted_rows = len(self.values) - self._replaced_rows
49
+
50
+
51
+ def insert_or_replace(
52
+ engine: sa.Engine,
53
+ table: sa.Table,
54
+ values: list[dict[str, T.Any]],
55
+ metadata: T.Optional[sa.MetaData] = None,
56
+ temp_table_name: T.Optional[str] = None,
57
+ conn: T.Optional[sa.Connection] = None,
58
+ trans: T.Optional[sa.Transaction] = None,
59
+ _raise_on_temp_table_create: bool = False,
60
+ _raise_on_temp_data_insert: bool = False,
61
+ _raise_on_target_delete: bool = False,
62
+ _raise_on_target_insert: bool = False,
63
+ _raise_on_temp_table_drop: bool = False,
64
+ ) -> tuple[int, int]:
65
+ """
66
+ Perform high-performance bulk INSERT-OR-REPLACE operation using temporary table.
67
+
68
+ This function performs bulk upsert operations: replaces existing records entirely
69
+ with new data and inserts records that don't exist. This is equivalent to
70
+ "INSERT OR REPLACE" or complete record replacement but works more efficiently
71
+ for large datasets.
72
+
73
+ **Algorithm**:
74
+
75
+ 1. Creates temporary table and loads all candidate data
76
+ 2. Uses JOIN to identify conflicting records in target table
77
+ 3. Deletes conflicting records from target table
78
+ 4. Bulk inserts all records from temporary table (both new and replacement)
79
+ 5. Cleans up temporary resources
80
+
81
+ This approach is ideal for:
82
+
83
+ - Full synchronization from authoritative data source
84
+ - Complete data refresh scenarios
85
+ - When new data should completely replace existing records
86
+
87
+ **Transaction Management**:
88
+
89
+ This function supports both auto-managed and user-managed transaction modes.
90
+ See the module-level documentation for detailed explanations of each mode.
91
+
92
+ :param engine: SQLAlchemy engine for database connection
93
+ :param table: Target table for upsert operation
94
+ :param values: Records to insert or replace.
95
+ Must include primary key values for conflict detection.
96
+ :param metadata: Optional metadata instance for temporary table isolation.
97
+ If None, a new MetaData instance is created for clean separation.
98
+ :param temp_table_name: Optional custom name for temporary table.
99
+ If None, generates unique name with timestamp to avoid conflicts.
100
+ :param conn: Optional database connection for user-managed transaction mode.
101
+ Must be provided together with ``trans`` parameter.
102
+ :param trans: Optional transaction for user-managed transaction mode.
103
+ Must be provided together with ``conn`` parameter.
104
+
105
+ :returns: Tuple of (replaced_rows, inserted_rows):
106
+ - replaced_rows: Number of existing records that were replaced
107
+ - inserted_rows: Number of new records that were inserted
108
+
109
+ :raises ValueError: When conn and trans parameters are provided inconsistently
110
+ (one is None while the other is not)
111
+ :raises UpsertTestError: When testing flags are enabled and corresponding operations fail
112
+
113
+ **Examples**:
114
+
115
+ Auto-managed transaction (default mode)::
116
+
117
+ # Function manages its own transaction
118
+ updated, inserted = insert_or_replace(engine, users_table, new_data)
119
+
120
+ User-managed transaction mode::
121
+
122
+ # Operation is part of larger transaction
123
+ with engine.connect() as conn:
124
+ with conn.begin() as trans:
125
+ # Other operations...
126
+ updated, inserted = insert_or_replace(
127
+ engine, users_table, new_data, conn=conn, trans=trans
128
+ )
129
+ # More operations...
130
+
131
+ Complete replacement example::
132
+
133
+ # Target table has records with id=1,2,3
134
+ new_data = [
135
+ {'id': 2, 'name': 'Bob Updated'}, # Exists - will be replaced
136
+ {'id': 4, 'name': 'Charlie'}, # New - will be inserted
137
+ {'id': 5, 'name': 'David'}, # New - will be inserted
138
+ ]
139
+ updated, inserted = insert_or_replace(engine, users_table, new_data)
140
+ # Result: updated=1, inserted=2
141
+
142
+ **Performance Comparison**:
143
+ Traditional row-by-row approach (100K records): ~300 seconds
144
+ This method (100K records): ~15 seconds
145
+ Performance gain: ~20x faster
146
+
147
+ .. note::
148
+
149
+ Parameters prefixed with ``_raise_on_`` are exclusively for testing error
150
+ handling and cleanup behavior. Never use these in production code.
151
+
152
+ .. warning::
153
+
154
+ This operation completely replaces existing records. All fields of
155
+ conflicting records (including historical fields like timestamps) will
156
+ be overwritten with new data.
157
+ """
158
+ if not values: # pragma: no cover
159
+ return 0, 0 # No-op for empty data
160
+
161
+ executor = InsertOrReplaceExcutor.new(
162
+ engine=engine,
163
+ table=table,
164
+ values=values,
165
+ metadata=metadata,
166
+ temp_table_name=temp_table_name,
167
+ conn=conn,
168
+ trans=trans,
169
+ _raise_on_temp_table_create=_raise_on_temp_table_create,
170
+ _raise_on_temp_data_insert=_raise_on_temp_data_insert,
171
+ _raise_on_target_delete=_raise_on_target_delete,
172
+ _raise_on_target_insert=_raise_on_target_insert,
173
+ _raise_on_temp_table_drop=_raise_on_temp_table_drop,
174
+ )
175
+ executor.run()
176
+ replaced_rows, inserted_rows = executor._replaced_rows, executor._inserted_rows
177
+ return replaced_rows, inserted_rows
@@ -0,0 +1,92 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ import typing as T
4
+ from datetime import datetime, timezone
5
+
6
+ import sqlalchemy as sa
7
+
8
+
9
+ def get_utc_now() -> datetime:
10
+ return datetime.now(tz=timezone.utc)
11
+
12
+
13
+ def get_pk_name(table: sa.Table) -> str:
14
+ """
15
+ Extract the primary key column name from a SQLAlchemy table.
16
+
17
+ This function ensures the table has exactly one primary key column,
18
+ which is required for the bulk operations to work correctly.
19
+
20
+ Args:
21
+ table: SQLAlchemy table object
22
+
23
+ Returns:
24
+ Name of the primary key column
25
+
26
+ Raises:
27
+ ValueError: If table has zero or multiple primary key columns
28
+
29
+ Example:
30
+ >>> table = sa.Table(
31
+ ... 'users',
32
+ ... metadata,
33
+ ... sa.Column('id', sa.Integer, primary_key=True)
34
+ ... )
35
+ >>> get_pk_name(table)
36
+ 'id'
37
+ """
38
+ pks = list(table.primary_key)
39
+ if len(pks) != 1: # pragma: no cover
40
+ pk_names = [pk.name for pk in pks]
41
+ raise ValueError(
42
+ f"Table must have exactly one primary key, but found: {pk_names}"
43
+ )
44
+ pk_name = pks[0].name
45
+ return pk_name
46
+
47
+
48
+ def get_temp_table_name(original_table_name: str) -> str:
49
+ dt = get_utc_now().strftime("%Y%m%d%H%M%S")
50
+ temp_table_name = f"temp_{dt}_" + original_table_name
51
+ return temp_table_name
52
+
53
+
54
+ def clone_temp_table(
55
+ original_table: sa.Table,
56
+ metadata: sa.MetaData,
57
+ temp_table_name: T.Optional[str] = None,
58
+ ) -> sa.Table:
59
+ """
60
+ Create a temporary table with the same schema as the original table.
61
+
62
+ This function clones the structure of an existing table to create a temporary
63
+ table for bulk operations. The temporary table inherits all columns, types,
64
+ and constraints from the original table.
65
+
66
+ Args:
67
+ original_table: The table to clone
68
+ metadata: Metadata object for the temporary table.
69
+ Should be a separate instance from the original table's metadata
70
+ to avoid conflicts and enable proper cleanup.
71
+ temp_table_name: Custom name for the temporary table.
72
+ If None, generates a unique name with timestamp to avoid conflicts
73
+ in concurrent environments.
74
+
75
+ Returns:
76
+ New temporary table with identical schema
77
+
78
+ .. note::
79
+
80
+ - Use a separate MetaData instance to isolate the temporary table
81
+ - In high-concurrency scenarios, consider providing unique temp_table_name
82
+ - The temporary table is not automatically bound to any engine
83
+
84
+ Example:
85
+ >>> metadata = sa.MetaData()
86
+ >>> _temp_table = clone_temp_table(users_table, metadata)
87
+ >>> # _temp_table has same columns as users_table but different metadata
88
+ """
89
+ if temp_table_name is None:
90
+ temp_table_name = get_temp_table_name(original_table.name)
91
+ temp_table = original_table.to_metadata(metadata, name=temp_table_name)
92
+ return temp_table
@@ -0,0 +1,2 @@
1
+ # -*- coding: utf-8 -*-
2
+