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.
- sqlalchemy_upsert_kit/__init__.py +9 -0
- sqlalchemy_upsert_kit/_version.py +15 -0
- sqlalchemy_upsert_kit/api.py +3 -0
- sqlalchemy_upsert_kit/exc.py +13 -0
- sqlalchemy_upsert_kit/paths.py +33 -0
- sqlalchemy_upsert_kit/sqlite/__init__.py +34 -0
- sqlalchemy_upsert_kit/sqlite/api.py +5 -0
- sqlalchemy_upsert_kit/sqlite/executor.py +535 -0
- sqlalchemy_upsert_kit/sqlite/insert_or_ignore.py +156 -0
- sqlalchemy_upsert_kit/sqlite/insert_or_merge.py +254 -0
- sqlalchemy_upsert_kit/sqlite/insert_or_replace.py +177 -0
- sqlalchemy_upsert_kit/utils.py +92 -0
- sqlalchemy_upsert_kit/vendor/__init__.py +2 -0
- sqlalchemy_upsert_kit/vendor/pytest_cov_helper.py +148 -0
- sqlalchemy_upsert_kit-0.1.1.dist-info/AUTHORS.rst +15 -0
- sqlalchemy_upsert_kit-0.1.1.dist-info/LICENSE.txt +21 -0
- sqlalchemy_upsert_kit-0.1.1.dist-info/METADATA +119 -0
- sqlalchemy_upsert_kit-0.1.1.dist-info/RECORD +19 -0
- sqlalchemy_upsert_kit-0.1.1.dist-info/WHEEL +4 -0
|
@@ -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
|