sqlspec 0.16.2__cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.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.

Potentially problematic release.


This version of sqlspec might be problematic. Click here for more details.

Files changed (148) hide show
  1. 51ff5a9eadfdefd49f98__mypyc.cpython-39-aarch64-linux-gnu.so +0 -0
  2. sqlspec/__init__.py +92 -0
  3. sqlspec/__main__.py +12 -0
  4. sqlspec/__metadata__.py +14 -0
  5. sqlspec/_serialization.py +77 -0
  6. sqlspec/_sql.py +1782 -0
  7. sqlspec/_typing.py +680 -0
  8. sqlspec/adapters/__init__.py +0 -0
  9. sqlspec/adapters/adbc/__init__.py +5 -0
  10. sqlspec/adapters/adbc/_types.py +12 -0
  11. sqlspec/adapters/adbc/config.py +361 -0
  12. sqlspec/adapters/adbc/driver.py +512 -0
  13. sqlspec/adapters/aiosqlite/__init__.py +19 -0
  14. sqlspec/adapters/aiosqlite/_types.py +13 -0
  15. sqlspec/adapters/aiosqlite/config.py +253 -0
  16. sqlspec/adapters/aiosqlite/driver.py +248 -0
  17. sqlspec/adapters/asyncmy/__init__.py +19 -0
  18. sqlspec/adapters/asyncmy/_types.py +12 -0
  19. sqlspec/adapters/asyncmy/config.py +180 -0
  20. sqlspec/adapters/asyncmy/driver.py +274 -0
  21. sqlspec/adapters/asyncpg/__init__.py +21 -0
  22. sqlspec/adapters/asyncpg/_types.py +17 -0
  23. sqlspec/adapters/asyncpg/config.py +229 -0
  24. sqlspec/adapters/asyncpg/driver.py +344 -0
  25. sqlspec/adapters/bigquery/__init__.py +18 -0
  26. sqlspec/adapters/bigquery/_types.py +12 -0
  27. sqlspec/adapters/bigquery/config.py +298 -0
  28. sqlspec/adapters/bigquery/driver.py +558 -0
  29. sqlspec/adapters/duckdb/__init__.py +22 -0
  30. sqlspec/adapters/duckdb/_types.py +12 -0
  31. sqlspec/adapters/duckdb/config.py +504 -0
  32. sqlspec/adapters/duckdb/driver.py +368 -0
  33. sqlspec/adapters/oracledb/__init__.py +32 -0
  34. sqlspec/adapters/oracledb/_types.py +14 -0
  35. sqlspec/adapters/oracledb/config.py +317 -0
  36. sqlspec/adapters/oracledb/driver.py +538 -0
  37. sqlspec/adapters/psqlpy/__init__.py +16 -0
  38. sqlspec/adapters/psqlpy/_types.py +11 -0
  39. sqlspec/adapters/psqlpy/config.py +214 -0
  40. sqlspec/adapters/psqlpy/driver.py +530 -0
  41. sqlspec/adapters/psycopg/__init__.py +32 -0
  42. sqlspec/adapters/psycopg/_types.py +17 -0
  43. sqlspec/adapters/psycopg/config.py +426 -0
  44. sqlspec/adapters/psycopg/driver.py +796 -0
  45. sqlspec/adapters/sqlite/__init__.py +15 -0
  46. sqlspec/adapters/sqlite/_types.py +11 -0
  47. sqlspec/adapters/sqlite/config.py +240 -0
  48. sqlspec/adapters/sqlite/driver.py +294 -0
  49. sqlspec/base.py +571 -0
  50. sqlspec/builder/__init__.py +62 -0
  51. sqlspec/builder/_base.py +473 -0
  52. sqlspec/builder/_column.py +320 -0
  53. sqlspec/builder/_ddl.py +1346 -0
  54. sqlspec/builder/_ddl_utils.py +103 -0
  55. sqlspec/builder/_delete.py +76 -0
  56. sqlspec/builder/_insert.py +421 -0
  57. sqlspec/builder/_merge.py +71 -0
  58. sqlspec/builder/_parsing_utils.py +164 -0
  59. sqlspec/builder/_select.py +170 -0
  60. sqlspec/builder/_update.py +188 -0
  61. sqlspec/builder/mixins/__init__.py +55 -0
  62. sqlspec/builder/mixins/_cte_and_set_ops.py +222 -0
  63. sqlspec/builder/mixins/_delete_operations.py +41 -0
  64. sqlspec/builder/mixins/_insert_operations.py +244 -0
  65. sqlspec/builder/mixins/_join_operations.py +149 -0
  66. sqlspec/builder/mixins/_merge_operations.py +562 -0
  67. sqlspec/builder/mixins/_order_limit_operations.py +135 -0
  68. sqlspec/builder/mixins/_pivot_operations.py +153 -0
  69. sqlspec/builder/mixins/_select_operations.py +604 -0
  70. sqlspec/builder/mixins/_update_operations.py +202 -0
  71. sqlspec/builder/mixins/_where_clause.py +644 -0
  72. sqlspec/cli.py +247 -0
  73. sqlspec/config.py +395 -0
  74. sqlspec/core/__init__.py +63 -0
  75. sqlspec/core/cache.cpython-39-aarch64-linux-gnu.so +0 -0
  76. sqlspec/core/cache.py +871 -0
  77. sqlspec/core/compiler.cpython-39-aarch64-linux-gnu.so +0 -0
  78. sqlspec/core/compiler.py +417 -0
  79. sqlspec/core/filters.cpython-39-aarch64-linux-gnu.so +0 -0
  80. sqlspec/core/filters.py +830 -0
  81. sqlspec/core/hashing.cpython-39-aarch64-linux-gnu.so +0 -0
  82. sqlspec/core/hashing.py +310 -0
  83. sqlspec/core/parameters.cpython-39-aarch64-linux-gnu.so +0 -0
  84. sqlspec/core/parameters.py +1237 -0
  85. sqlspec/core/result.cpython-39-aarch64-linux-gnu.so +0 -0
  86. sqlspec/core/result.py +677 -0
  87. sqlspec/core/splitter.cpython-39-aarch64-linux-gnu.so +0 -0
  88. sqlspec/core/splitter.py +819 -0
  89. sqlspec/core/statement.cpython-39-aarch64-linux-gnu.so +0 -0
  90. sqlspec/core/statement.py +676 -0
  91. sqlspec/driver/__init__.py +19 -0
  92. sqlspec/driver/_async.py +502 -0
  93. sqlspec/driver/_common.py +631 -0
  94. sqlspec/driver/_sync.py +503 -0
  95. sqlspec/driver/mixins/__init__.py +6 -0
  96. sqlspec/driver/mixins/_result_tools.py +193 -0
  97. sqlspec/driver/mixins/_sql_translator.py +86 -0
  98. sqlspec/exceptions.py +193 -0
  99. sqlspec/extensions/__init__.py +0 -0
  100. sqlspec/extensions/aiosql/__init__.py +10 -0
  101. sqlspec/extensions/aiosql/adapter.py +461 -0
  102. sqlspec/extensions/litestar/__init__.py +6 -0
  103. sqlspec/extensions/litestar/_utils.py +52 -0
  104. sqlspec/extensions/litestar/cli.py +48 -0
  105. sqlspec/extensions/litestar/config.py +92 -0
  106. sqlspec/extensions/litestar/handlers.py +260 -0
  107. sqlspec/extensions/litestar/plugin.py +145 -0
  108. sqlspec/extensions/litestar/providers.py +454 -0
  109. sqlspec/loader.cpython-39-aarch64-linux-gnu.so +0 -0
  110. sqlspec/loader.py +760 -0
  111. sqlspec/migrations/__init__.py +35 -0
  112. sqlspec/migrations/base.py +414 -0
  113. sqlspec/migrations/commands.py +443 -0
  114. sqlspec/migrations/loaders.py +402 -0
  115. sqlspec/migrations/runner.py +213 -0
  116. sqlspec/migrations/tracker.py +140 -0
  117. sqlspec/migrations/utils.py +129 -0
  118. sqlspec/protocols.py +407 -0
  119. sqlspec/py.typed +0 -0
  120. sqlspec/storage/__init__.py +23 -0
  121. sqlspec/storage/backends/__init__.py +0 -0
  122. sqlspec/storage/backends/base.py +163 -0
  123. sqlspec/storage/backends/fsspec.py +386 -0
  124. sqlspec/storage/backends/obstore.py +459 -0
  125. sqlspec/storage/capabilities.py +102 -0
  126. sqlspec/storage/registry.py +239 -0
  127. sqlspec/typing.py +299 -0
  128. sqlspec/utils/__init__.py +3 -0
  129. sqlspec/utils/correlation.py +150 -0
  130. sqlspec/utils/deprecation.py +106 -0
  131. sqlspec/utils/fixtures.cpython-39-aarch64-linux-gnu.so +0 -0
  132. sqlspec/utils/fixtures.py +58 -0
  133. sqlspec/utils/logging.py +127 -0
  134. sqlspec/utils/module_loader.py +89 -0
  135. sqlspec/utils/serializers.py +4 -0
  136. sqlspec/utils/singleton.py +32 -0
  137. sqlspec/utils/sync_tools.cpython-39-aarch64-linux-gnu.so +0 -0
  138. sqlspec/utils/sync_tools.py +237 -0
  139. sqlspec/utils/text.cpython-39-aarch64-linux-gnu.so +0 -0
  140. sqlspec/utils/text.py +96 -0
  141. sqlspec/utils/type_guards.cpython-39-aarch64-linux-gnu.so +0 -0
  142. sqlspec/utils/type_guards.py +1139 -0
  143. sqlspec-0.16.2.dist-info/METADATA +365 -0
  144. sqlspec-0.16.2.dist-info/RECORD +148 -0
  145. sqlspec-0.16.2.dist-info/WHEEL +7 -0
  146. sqlspec-0.16.2.dist-info/entry_points.txt +2 -0
  147. sqlspec-0.16.2.dist-info/licenses/LICENSE +21 -0
  148. sqlspec-0.16.2.dist-info/licenses/NOTICE +29 -0
@@ -0,0 +1,562 @@
1
+ """Merge operation mixins for SQL builders."""
2
+
3
+ from typing import Any, Optional, Union
4
+
5
+ from mypy_extensions import trait
6
+ from sqlglot import exp
7
+ from typing_extensions import Self
8
+
9
+ from sqlspec.exceptions import SQLBuilderError
10
+ from sqlspec.utils.type_guards import has_query_builder_parameters
11
+
12
+ __all__ = (
13
+ "MergeIntoClauseMixin",
14
+ "MergeMatchedClauseMixin",
15
+ "MergeNotMatchedBySourceClauseMixin",
16
+ "MergeNotMatchedClauseMixin",
17
+ "MergeOnClauseMixin",
18
+ "MergeUsingClauseMixin",
19
+ )
20
+
21
+
22
+ @trait
23
+ class MergeIntoClauseMixin:
24
+ """Mixin providing INTO clause for MERGE builders."""
25
+
26
+ __slots__ = ()
27
+ _expression: Optional[exp.Expression]
28
+
29
+ def into(self, table: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
30
+ """Set the target table for the MERGE operation (INTO clause).
31
+
32
+ Args:
33
+ table: The target table name or expression for the MERGE operation.
34
+ Can be a string (table name) or an sqlglot Expression.
35
+ alias: Optional alias for the target table.
36
+
37
+ Returns:
38
+ The current builder instance for method chaining.
39
+ """
40
+ if self._expression is None:
41
+ self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
42
+ if not isinstance(self._expression, exp.Merge):
43
+ self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
44
+ self._expression.set("this", exp.to_table(table, alias=alias) if isinstance(table, str) else table)
45
+ return self
46
+
47
+
48
+ @trait
49
+ class MergeUsingClauseMixin:
50
+ """Mixin providing USING clause for MERGE builders."""
51
+
52
+ __slots__ = ()
53
+ _expression: Optional[exp.Expression]
54
+
55
+ def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
56
+ """Add parameter - provided by QueryBuilder."""
57
+ msg = "Method must be provided by QueryBuilder subclass"
58
+ raise NotImplementedError(msg)
59
+
60
+ def using(self, source: Union[str, exp.Expression, Any], alias: Optional[str] = None) -> Self:
61
+ """Set the source data for the MERGE operation (USING clause).
62
+
63
+ Args:
64
+ source: The source data for the MERGE operation.
65
+ Can be a string (table name), an sqlglot Expression, or a SelectBuilder instance.
66
+ alias: Optional alias for the source table.
67
+
68
+ Returns:
69
+ The current builder instance for method chaining.
70
+
71
+ Raises:
72
+ SQLBuilderError: If the current expression is not a MERGE statement or if the source type is unsupported.
73
+ """
74
+ if self._expression is None:
75
+ self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
76
+ if not isinstance(self._expression, exp.Merge):
77
+ self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
78
+
79
+ source_expr: exp.Expression
80
+ if isinstance(source, str):
81
+ source_expr = exp.to_table(source, alias=alias)
82
+ elif has_query_builder_parameters(source) and hasattr(source, "_expression"):
83
+ subquery_builder_parameters = source.parameters
84
+ if subquery_builder_parameters:
85
+ for p_name, p_value in subquery_builder_parameters.items():
86
+ self.add_parameter(p_value, name=p_name)
87
+
88
+ subquery_exp = exp.paren(getattr(source, "_expression", exp.select()))
89
+ source_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
90
+ elif isinstance(source, exp.Expression):
91
+ source_expr = source
92
+ if alias:
93
+ source_expr = exp.alias_(source_expr, alias)
94
+ else:
95
+ msg = f"Unsupported source type for USING clause: {type(source)}"
96
+ raise SQLBuilderError(msg)
97
+
98
+ self._expression.set("using", source_expr)
99
+ return self
100
+
101
+
102
+ @trait
103
+ class MergeOnClauseMixin:
104
+ """Mixin providing ON clause for MERGE builders."""
105
+
106
+ __slots__ = ()
107
+ _expression: Optional[exp.Expression]
108
+
109
+ def on(self, condition: Union[str, exp.Expression]) -> Self:
110
+ """Set the join condition for the MERGE operation (ON clause).
111
+
112
+ Args:
113
+ condition: The join condition for the MERGE operation.
114
+ Can be a string (SQL condition) or an sqlglot Expression.
115
+
116
+ Returns:
117
+ The current builder instance for method chaining.
118
+
119
+ Raises:
120
+ SQLBuilderError: If the current expression is not a MERGE statement or if the condition type is unsupported.
121
+ """
122
+ if self._expression is None:
123
+ self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
124
+ if not isinstance(self._expression, exp.Merge):
125
+ self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
126
+
127
+ condition_expr: exp.Expression
128
+ if isinstance(condition, str):
129
+ parsed_condition: Optional[exp.Expression] = exp.maybe_parse(
130
+ condition, dialect=getattr(self, "dialect", None)
131
+ )
132
+ if not parsed_condition:
133
+ msg = f"Could not parse ON condition: {condition}"
134
+ raise SQLBuilderError(msg)
135
+ condition_expr = parsed_condition
136
+ elif isinstance(condition, exp.Expression):
137
+ condition_expr = condition
138
+ else:
139
+ msg = f"Unsupported condition type for ON clause: {type(condition)}"
140
+ raise SQLBuilderError(msg)
141
+
142
+ self._expression.set("on", condition_expr)
143
+ return self
144
+
145
+
146
+ @trait
147
+ class MergeMatchedClauseMixin:
148
+ """Mixin providing WHEN MATCHED THEN ... clauses for MERGE builders."""
149
+
150
+ __slots__ = ()
151
+ _expression: Optional[exp.Expression]
152
+
153
+ def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
154
+ """Add parameter - provided by QueryBuilder."""
155
+ msg = "Method must be provided by QueryBuilder subclass"
156
+ raise NotImplementedError(msg)
157
+
158
+ def _generate_unique_parameter_name(self, base_name: str) -> str:
159
+ """Generate unique parameter name - provided by QueryBuilder."""
160
+ msg = "Method must be provided by QueryBuilder subclass"
161
+ raise NotImplementedError(msg)
162
+
163
+ def _add_when_clause(self, when_clause: exp.When) -> None:
164
+ """Helper to add a WHEN clause to the MERGE statement.
165
+
166
+ Args:
167
+ when_clause: The WHEN clause to add.
168
+ """
169
+ if self._expression is None:
170
+ self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[])) # type: ignore[misc]
171
+ if not isinstance(self._expression, exp.Merge):
172
+ self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[])) # type: ignore[misc]
173
+
174
+ whens = self._expression.args.get("whens")
175
+ if not whens:
176
+ whens = exp.Whens(expressions=[])
177
+ self._expression.set("whens", whens)
178
+
179
+ whens.append("expressions", when_clause)
180
+
181
+ def when_matched_then_update(
182
+ self,
183
+ set_values: Optional[dict[str, Any]] = None,
184
+ condition: Optional[Union[str, exp.Expression]] = None,
185
+ **kwargs: Any,
186
+ ) -> Self:
187
+ """Define the UPDATE action for matched rows.
188
+
189
+ Supports:
190
+ - when_matched_then_update({"column": value})
191
+ - when_matched_then_update(column=value, other_column=other_value)
192
+ - when_matched_then_update({"column": value}, other_column=other_value)
193
+
194
+ Args:
195
+ set_values: A dictionary of column names and their new values to set.
196
+ The values will be parameterized.
197
+ condition: An optional additional condition for this specific action.
198
+ **kwargs: Column-value pairs to update on match.
199
+
200
+ Raises:
201
+ SQLBuilderError: If the condition type is unsupported.
202
+
203
+ Returns:
204
+ The current builder instance for method chaining.
205
+ """
206
+ # Combine set_values dict and kwargs
207
+ all_values = dict(set_values or {}, **kwargs)
208
+
209
+ if not all_values:
210
+ msg = "No update values provided. Use set_values dict or kwargs."
211
+ raise SQLBuilderError(msg)
212
+
213
+ update_expressions: list[exp.EQ] = []
214
+ for col, val in all_values.items():
215
+ if hasattr(val, "expression") and hasattr(val, "sql"):
216
+ # Handle SQL objects (from sql.raw with parameters)
217
+ expression = getattr(val, "expression", None)
218
+ if expression is not None and isinstance(expression, exp.Expression):
219
+ # Merge parameters from SQL object into builder
220
+ if hasattr(val, "parameters"):
221
+ sql_parameters = getattr(val, "parameters", {})
222
+ for param_name, param_value in sql_parameters.items():
223
+ self.add_parameter(param_value, name=param_name)
224
+ value_expr = expression
225
+ else:
226
+ # If expression is None, fall back to parsing the raw SQL
227
+ sql_text = getattr(val, "sql", "")
228
+ # Merge parameters even when parsing raw SQL
229
+ if hasattr(val, "parameters"):
230
+ sql_parameters = getattr(val, "parameters", {})
231
+ for param_name, param_value in sql_parameters.items():
232
+ self.add_parameter(param_value, name=param_name)
233
+ # Check if sql_text is callable (like Expression.sql method)
234
+ if callable(sql_text):
235
+ sql_text = str(val)
236
+ value_expr = exp.maybe_parse(sql_text) or exp.convert(str(sql_text))
237
+ elif isinstance(val, exp.Expression):
238
+ value_expr = val
239
+ else:
240
+ column_name = col if isinstance(col, str) else str(col)
241
+ if "." in column_name:
242
+ column_name = column_name.split(".")[-1]
243
+ param_name = self._generate_unique_parameter_name(column_name)
244
+ param_name = self.add_parameter(val, name=param_name)[1]
245
+ value_expr = exp.Placeholder(this=param_name)
246
+
247
+ update_expressions.append(exp.EQ(this=exp.column(col), expression=value_expr))
248
+
249
+ when_args: dict[str, Any] = {"matched": True, "then": exp.Update(expressions=update_expressions)}
250
+
251
+ if condition:
252
+ condition_expr: exp.Expression
253
+ if isinstance(condition, str):
254
+ parsed_cond: Optional[exp.Expression] = exp.maybe_parse(
255
+ condition, dialect=getattr(self, "dialect", None)
256
+ )
257
+ if not parsed_cond:
258
+ msg = f"Could not parse WHEN clause condition: {condition}"
259
+ raise SQLBuilderError(msg)
260
+ condition_expr = parsed_cond
261
+ elif isinstance(condition, exp.Expression):
262
+ condition_expr = condition
263
+ else:
264
+ msg = f"Unsupported condition type for WHEN clause: {type(condition)}"
265
+ raise SQLBuilderError(msg)
266
+ when_args["this"] = condition_expr
267
+
268
+ when_clause = exp.When(**when_args)
269
+ self._add_when_clause(when_clause)
270
+ return self
271
+
272
+ def when_matched_then_delete(self, condition: Optional[Union[str, exp.Expression]] = None) -> Self:
273
+ """Define the DELETE action for matched rows.
274
+
275
+ Args:
276
+ condition: An optional additional condition for this specific action.
277
+
278
+ Raises:
279
+ SQLBuilderError: If the condition type is unsupported.
280
+
281
+ Returns:
282
+ The current builder instance for method chaining.
283
+ """
284
+ when_args: dict[str, Any] = {"matched": True, "then": exp.Delete()}
285
+
286
+ if condition:
287
+ condition_expr: exp.Expression
288
+ if isinstance(condition, str):
289
+ parsed_cond: Optional[exp.Expression] = exp.maybe_parse(
290
+ condition, dialect=getattr(self, "dialect", None)
291
+ )
292
+ if not parsed_cond:
293
+ msg = f"Could not parse WHEN clause condition: {condition}"
294
+ raise SQLBuilderError(msg)
295
+ condition_expr = parsed_cond
296
+ elif isinstance(condition, exp.Expression):
297
+ condition_expr = condition
298
+ else:
299
+ msg = f"Unsupported condition type for WHEN clause: {type(condition)}"
300
+ raise SQLBuilderError(msg)
301
+ when_args["this"] = condition_expr
302
+
303
+ when_clause = exp.When(**when_args)
304
+ self._add_when_clause(when_clause)
305
+ return self
306
+
307
+
308
+ @trait
309
+ class MergeNotMatchedClauseMixin:
310
+ """Mixin providing WHEN NOT MATCHED THEN ... clauses for MERGE builders."""
311
+
312
+ __slots__ = ()
313
+
314
+ _expression: Optional[exp.Expression]
315
+
316
+ def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
317
+ """Add parameter - provided by QueryBuilder."""
318
+ msg = "Method must be provided by QueryBuilder subclass"
319
+ raise NotImplementedError(msg)
320
+
321
+ def _generate_unique_parameter_name(self, base_name: str) -> str:
322
+ """Generate unique parameter name - provided by QueryBuilder."""
323
+ msg = "Method must be provided by QueryBuilder subclass"
324
+ raise NotImplementedError(msg)
325
+
326
+ def _add_when_clause(self, when_clause: exp.When) -> None:
327
+ """Helper to add a WHEN clause to the MERGE statement - provided by QueryBuilder."""
328
+ msg = "Method must be provided by QueryBuilder subclass"
329
+ raise NotImplementedError(msg)
330
+
331
+ def when_not_matched_then_insert(
332
+ self,
333
+ columns: Optional[list[str]] = None,
334
+ values: Optional[list[Any]] = None,
335
+ condition: Optional[Union[str, exp.Expression]] = None,
336
+ by_target: bool = True,
337
+ ) -> Self:
338
+ """Define the INSERT action for rows not matched.
339
+
340
+ Args:
341
+ columns: A list of column names to insert into. If None, implies INSERT DEFAULT VALUES or matching source columns.
342
+ values: A list of values corresponding to the columns.
343
+ These values will be parameterized. If None, implies INSERT DEFAULT VALUES or subquery source.
344
+ condition: An optional additional condition for this specific action.
345
+ by_target: If True (default), condition is "WHEN NOT MATCHED [BY TARGET]".
346
+ If False, condition is "WHEN NOT MATCHED BY SOURCE".
347
+
348
+ Returns:
349
+ The current builder instance for method chaining.
350
+
351
+ Raises:
352
+ SQLBuilderError: If columns and values are provided but do not match in length,
353
+ or if columns are provided without values.
354
+ """
355
+ insert_args: dict[str, Any] = {}
356
+ if columns and values:
357
+ if len(columns) != len(values):
358
+ msg = "Number of columns must match number of values for INSERT."
359
+ raise SQLBuilderError(msg)
360
+
361
+ parameterized_values: list[exp.Expression] = []
362
+ for i, val in enumerate(values):
363
+ column_name = columns[i] if isinstance(columns[i], str) else str(columns[i])
364
+ if "." in column_name:
365
+ column_name = column_name.split(".")[-1]
366
+ param_name = self._generate_unique_parameter_name(column_name)
367
+ param_name = self.add_parameter(val, name=param_name)[1]
368
+ parameterized_values.append(exp.var(param_name))
369
+
370
+ insert_args["this"] = exp.Tuple(expressions=[exp.column(c) for c in columns])
371
+ insert_args["expression"] = exp.Tuple(expressions=parameterized_values)
372
+ elif columns and not values:
373
+ msg = "Specifying columns without values for INSERT action is complex and not fully supported yet. Consider providing full expressions."
374
+ raise SQLBuilderError(msg)
375
+ elif not columns and not values:
376
+ pass
377
+ else:
378
+ msg = "Cannot specify values without columns for INSERT action."
379
+ raise SQLBuilderError(msg)
380
+
381
+ when_args: dict[str, Any] = {"matched": False, "then": exp.Insert(**insert_args)}
382
+
383
+ if not by_target:
384
+ when_args["source"] = True
385
+
386
+ if condition:
387
+ condition_expr: exp.Expression
388
+ if isinstance(condition, str):
389
+ parsed_cond: Optional[exp.Expression] = exp.maybe_parse(
390
+ condition, dialect=getattr(self, "dialect", None)
391
+ )
392
+ if not parsed_cond:
393
+ msg = f"Could not parse WHEN clause condition: {condition}"
394
+ raise SQLBuilderError(msg)
395
+ condition_expr = parsed_cond
396
+ elif isinstance(condition, exp.Expression):
397
+ condition_expr = condition
398
+ else:
399
+ msg = f"Unsupported condition type for WHEN clause: {type(condition)}"
400
+ raise SQLBuilderError(msg)
401
+ when_args["this"] = condition_expr
402
+
403
+ when_clause = exp.When(**when_args)
404
+ self._add_when_clause(when_clause)
405
+ return self
406
+
407
+
408
+ @trait
409
+ class MergeNotMatchedBySourceClauseMixin:
410
+ """Mixin providing WHEN NOT MATCHED BY SOURCE THEN ... clauses for MERGE builders."""
411
+
412
+ __slots__ = ()
413
+
414
+ _expression: Optional[exp.Expression]
415
+
416
+ def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
417
+ """Add parameter - provided by QueryBuilder."""
418
+ msg = "Method must be provided by QueryBuilder subclass"
419
+ raise NotImplementedError(msg)
420
+
421
+ def _generate_unique_parameter_name(self, base_name: str) -> str:
422
+ """Generate unique parameter name - provided by QueryBuilder."""
423
+ msg = "Method must be provided by QueryBuilder subclass"
424
+ raise NotImplementedError(msg)
425
+
426
+ def _add_when_clause(self, when_clause: exp.When) -> None:
427
+ """Helper to add a WHEN clause to the MERGE statement - provided by QueryBuilder."""
428
+ msg = "Method must be provided by QueryBuilder subclass"
429
+ raise NotImplementedError(msg)
430
+
431
+ def when_not_matched_by_source_then_update(
432
+ self,
433
+ set_values: Optional[dict[str, Any]] = None,
434
+ condition: Optional[Union[str, exp.Expression]] = None,
435
+ **kwargs: Any,
436
+ ) -> Self:
437
+ """Define the UPDATE action for rows not matched by source.
438
+
439
+ This is useful for handling rows that exist in the target but not in the source.
440
+
441
+ Supports:
442
+ - when_not_matched_by_source_then_update({"column": value})
443
+ - when_not_matched_by_source_then_update(column=value, other_column=other_value)
444
+ - when_not_matched_by_source_then_update({"column": value}, other_column=other_value)
445
+
446
+ Args:
447
+ set_values: A dictionary of column names and their new values to set.
448
+ condition: An optional additional condition for this specific action.
449
+ **kwargs: Column-value pairs to update when not matched by source.
450
+
451
+ Raises:
452
+ SQLBuilderError: If the condition type is unsupported.
453
+
454
+ Returns:
455
+ The current builder instance for method chaining.
456
+ """
457
+ # Combine set_values dict and kwargs
458
+ all_values = dict(set_values or {}, **kwargs)
459
+
460
+ if not all_values:
461
+ msg = "No update values provided. Use set_values dict or kwargs."
462
+ raise SQLBuilderError(msg)
463
+
464
+ update_expressions: list[exp.EQ] = []
465
+ for col, val in all_values.items():
466
+ if hasattr(val, "expression") and hasattr(val, "sql"):
467
+ # Handle SQL objects (from sql.raw with parameters)
468
+ expression = getattr(val, "expression", None)
469
+ if expression is not None and isinstance(expression, exp.Expression):
470
+ # Merge parameters from SQL object into builder
471
+ if hasattr(val, "parameters"):
472
+ sql_parameters = getattr(val, "parameters", {})
473
+ for param_name, param_value in sql_parameters.items():
474
+ self.add_parameter(param_value, name=param_name)
475
+ value_expr = expression
476
+ else:
477
+ # If expression is None, fall back to parsing the raw SQL
478
+ sql_text = getattr(val, "sql", "")
479
+ # Merge parameters even when parsing raw SQL
480
+ if hasattr(val, "parameters"):
481
+ sql_parameters = getattr(val, "parameters", {})
482
+ for param_name, param_value in sql_parameters.items():
483
+ self.add_parameter(param_value, name=param_name)
484
+ # Check if sql_text is callable (like Expression.sql method)
485
+ if callable(sql_text):
486
+ sql_text = str(val)
487
+ value_expr = exp.maybe_parse(sql_text) or exp.convert(str(sql_text))
488
+ elif isinstance(val, exp.Expression):
489
+ value_expr = val
490
+ else:
491
+ column_name = col if isinstance(col, str) else str(col)
492
+ if "." in column_name:
493
+ column_name = column_name.split(".")[-1]
494
+ param_name = self._generate_unique_parameter_name(column_name)
495
+ param_name = self.add_parameter(val, name=param_name)[1]
496
+ value_expr = exp.Placeholder(this=param_name)
497
+
498
+ update_expressions.append(exp.EQ(this=exp.column(col), expression=value_expr))
499
+
500
+ when_args: dict[str, Any] = {
501
+ "matched": False,
502
+ "source": True,
503
+ "then": exp.Update(expressions=update_expressions),
504
+ }
505
+
506
+ if condition:
507
+ condition_expr: exp.Expression
508
+ if isinstance(condition, str):
509
+ parsed_cond: Optional[exp.Expression] = exp.maybe_parse(
510
+ condition, dialect=getattr(self, "dialect", None)
511
+ )
512
+ if not parsed_cond:
513
+ msg = f"Could not parse WHEN clause condition: {condition}"
514
+ raise SQLBuilderError(msg)
515
+ condition_expr = parsed_cond
516
+ elif isinstance(condition, exp.Expression):
517
+ condition_expr = condition
518
+ else:
519
+ msg = f"Unsupported condition type for WHEN clause: {type(condition)}"
520
+ raise SQLBuilderError(msg)
521
+ when_args["this"] = condition_expr
522
+
523
+ when_clause = exp.When(**when_args)
524
+ self._add_when_clause(when_clause)
525
+ return self
526
+
527
+ def when_not_matched_by_source_then_delete(self, condition: Optional[Union[str, exp.Expression]] = None) -> Self:
528
+ """Define the DELETE action for rows not matched by source.
529
+
530
+ This is useful for cleaning up rows that exist in the target but not in the source.
531
+
532
+ Args:
533
+ condition: An optional additional condition for this specific action.
534
+
535
+ Raises:
536
+ SQLBuilderError: If the condition type is unsupported.
537
+
538
+ Returns:
539
+ The current builder instance for method chaining.
540
+ """
541
+ when_args: dict[str, Any] = {"matched": False, "source": True, "then": exp.Delete()}
542
+
543
+ if condition:
544
+ condition_expr: exp.Expression
545
+ if isinstance(condition, str):
546
+ parsed_cond: Optional[exp.Expression] = exp.maybe_parse(
547
+ condition, dialect=getattr(self, "dialect", None)
548
+ )
549
+ if not parsed_cond:
550
+ msg = f"Could not parse WHEN clause condition: {condition}"
551
+ raise SQLBuilderError(msg)
552
+ condition_expr = parsed_cond
553
+ elif isinstance(condition, exp.Expression):
554
+ condition_expr = condition
555
+ else:
556
+ msg = f"Unsupported condition type for WHEN clause: {type(condition)}"
557
+ raise SQLBuilderError(msg)
558
+ when_args["this"] = condition_expr
559
+
560
+ when_clause = exp.When(**when_args)
561
+ self._add_when_clause(when_clause)
562
+ return self
@@ -0,0 +1,135 @@
1
+ """Order, Limit, Offset and Returning operations mixins for SQL builders."""
2
+
3
+ from typing import TYPE_CHECKING, Optional, Union, cast
4
+
5
+ from mypy_extensions import trait
6
+ from sqlglot import exp
7
+ from typing_extensions import Self
8
+
9
+ from sqlspec.builder._parsing_utils import parse_order_expression
10
+ from sqlspec.exceptions import SQLBuilderError
11
+
12
+ if TYPE_CHECKING:
13
+ from sqlspec.protocols import SQLBuilderProtocol
14
+
15
+ __all__ = ("LimitOffsetClauseMixin", "OrderByClauseMixin", "ReturningClauseMixin")
16
+
17
+
18
+ @trait
19
+ class OrderByClauseMixin:
20
+ """Mixin providing ORDER BY clause."""
21
+
22
+ __slots__ = ()
23
+
24
+ # Type annotation for PyRight - this will be provided by the base class
25
+ _expression: Optional[exp.Expression]
26
+
27
+ def order_by(self, *items: Union[str, exp.Ordered], desc: bool = False) -> Self:
28
+ """Add ORDER BY clause.
29
+
30
+ Args:
31
+ *items: Columns to order by. Can be strings (column names) or sqlglot.exp.Ordered instances for specific directions (e.g., exp.column("name").desc()).
32
+ desc: Whether to order in descending order (applies to all items if they are strings).
33
+
34
+ Raises:
35
+ SQLBuilderError: If the current expression is not a SELECT statement or if the item type is unsupported.
36
+
37
+ Returns:
38
+ The current builder instance for method chaining.
39
+ """
40
+ builder = cast("SQLBuilderProtocol", self)
41
+ if not isinstance(builder._expression, exp.Select):
42
+ msg = "ORDER BY is only supported for SELECT statements."
43
+ raise SQLBuilderError(msg)
44
+
45
+ current_expr = builder._expression
46
+ for item in items:
47
+ if isinstance(item, str):
48
+ order_item = parse_order_expression(item)
49
+ if desc:
50
+ order_item = order_item.desc()
51
+ else:
52
+ order_item = item
53
+ current_expr = current_expr.order_by(order_item, copy=False)
54
+ builder._expression = current_expr
55
+ return cast("Self", builder)
56
+
57
+
58
+ @trait
59
+ class LimitOffsetClauseMixin:
60
+ """Mixin providing LIMIT and OFFSET clauses."""
61
+
62
+ __slots__ = ()
63
+
64
+ # Type annotation for PyRight - this will be provided by the base class
65
+ _expression: Optional[exp.Expression]
66
+
67
+ def limit(self, value: int) -> Self:
68
+ """Add LIMIT clause.
69
+
70
+ Args:
71
+ value: The maximum number of rows to return.
72
+
73
+ Raises:
74
+ SQLBuilderError: If the current expression is not a SELECT statement.
75
+
76
+ Returns:
77
+ The current builder instance for method chaining.
78
+ """
79
+ builder = cast("SQLBuilderProtocol", self)
80
+ if not isinstance(builder._expression, exp.Select):
81
+ msg = "LIMIT is only supported for SELECT statements."
82
+ raise SQLBuilderError(msg)
83
+ builder._expression = builder._expression.limit(exp.convert(value), copy=False)
84
+ return cast("Self", builder)
85
+
86
+ def offset(self, value: int) -> Self:
87
+ """Add OFFSET clause.
88
+
89
+ Args:
90
+ value: The number of rows to skip before starting to return rows.
91
+
92
+ Raises:
93
+ SQLBuilderError: If the current expression is not a SELECT statement.
94
+
95
+ Returns:
96
+ The current builder instance for method chaining.
97
+ """
98
+ builder = cast("SQLBuilderProtocol", self)
99
+ if not isinstance(builder._expression, exp.Select):
100
+ msg = "OFFSET is only supported for SELECT statements."
101
+ raise SQLBuilderError(msg)
102
+ builder._expression = builder._expression.offset(exp.convert(value), copy=False)
103
+ return cast("Self", builder)
104
+
105
+
106
+ @trait
107
+ class ReturningClauseMixin:
108
+ """Mixin providing RETURNING clause."""
109
+
110
+ __slots__ = ()
111
+ # Type annotation for PyRight - this will be provided by the base class
112
+ _expression: Optional[exp.Expression]
113
+
114
+ def returning(self, *columns: Union[str, exp.Expression]) -> Self:
115
+ """Add RETURNING clause to the statement.
116
+
117
+ Args:
118
+ *columns: Columns to return. Can be strings or sqlglot expressions.
119
+
120
+ Raises:
121
+ SQLBuilderError: If the current expression is not INSERT, UPDATE, or DELETE.
122
+
123
+ Returns:
124
+ The current builder instance for method chaining.
125
+ """
126
+ if self._expression is None:
127
+ msg = "Cannot add RETURNING: expression is not initialized."
128
+ raise SQLBuilderError(msg)
129
+ valid_types = (exp.Insert, exp.Update, exp.Delete)
130
+ if not isinstance(self._expression, valid_types):
131
+ msg = "RETURNING is only supported for INSERT, UPDATE, and DELETE statements."
132
+ raise SQLBuilderError(msg)
133
+ returning_exprs = [exp.column(c) if isinstance(c, str) else c for c in columns]
134
+ self._expression.set("returning", exp.Returning(expressions=returning_exprs))
135
+ return self