sqlspec 0.26.0__py3-none-any.whl → 0.28.0__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.

Potentially problematic release.


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

Files changed (212) hide show
  1. sqlspec/__init__.py +7 -15
  2. sqlspec/_serialization.py +55 -25
  3. sqlspec/_typing.py +155 -52
  4. sqlspec/adapters/adbc/_types.py +1 -1
  5. sqlspec/adapters/adbc/adk/__init__.py +5 -0
  6. sqlspec/adapters/adbc/adk/store.py +880 -0
  7. sqlspec/adapters/adbc/config.py +62 -12
  8. sqlspec/adapters/adbc/data_dictionary.py +74 -2
  9. sqlspec/adapters/adbc/driver.py +226 -58
  10. sqlspec/adapters/adbc/litestar/__init__.py +5 -0
  11. sqlspec/adapters/adbc/litestar/store.py +504 -0
  12. sqlspec/adapters/adbc/type_converter.py +44 -50
  13. sqlspec/adapters/aiosqlite/_types.py +1 -1
  14. sqlspec/adapters/aiosqlite/adk/__init__.py +5 -0
  15. sqlspec/adapters/aiosqlite/adk/store.py +536 -0
  16. sqlspec/adapters/aiosqlite/config.py +86 -16
  17. sqlspec/adapters/aiosqlite/data_dictionary.py +34 -2
  18. sqlspec/adapters/aiosqlite/driver.py +127 -38
  19. sqlspec/adapters/aiosqlite/litestar/__init__.py +5 -0
  20. sqlspec/adapters/aiosqlite/litestar/store.py +281 -0
  21. sqlspec/adapters/aiosqlite/pool.py +7 -7
  22. sqlspec/adapters/asyncmy/__init__.py +7 -1
  23. sqlspec/adapters/asyncmy/_types.py +1 -1
  24. sqlspec/adapters/asyncmy/adk/__init__.py +5 -0
  25. sqlspec/adapters/asyncmy/adk/store.py +503 -0
  26. sqlspec/adapters/asyncmy/config.py +59 -17
  27. sqlspec/adapters/asyncmy/data_dictionary.py +41 -2
  28. sqlspec/adapters/asyncmy/driver.py +293 -62
  29. sqlspec/adapters/asyncmy/litestar/__init__.py +5 -0
  30. sqlspec/adapters/asyncmy/litestar/store.py +296 -0
  31. sqlspec/adapters/asyncpg/__init__.py +2 -1
  32. sqlspec/adapters/asyncpg/_type_handlers.py +71 -0
  33. sqlspec/adapters/asyncpg/_types.py +11 -7
  34. sqlspec/adapters/asyncpg/adk/__init__.py +5 -0
  35. sqlspec/adapters/asyncpg/adk/store.py +460 -0
  36. sqlspec/adapters/asyncpg/config.py +57 -36
  37. sqlspec/adapters/asyncpg/data_dictionary.py +48 -2
  38. sqlspec/adapters/asyncpg/driver.py +153 -23
  39. sqlspec/adapters/asyncpg/litestar/__init__.py +5 -0
  40. sqlspec/adapters/asyncpg/litestar/store.py +253 -0
  41. sqlspec/adapters/bigquery/_types.py +1 -1
  42. sqlspec/adapters/bigquery/adk/__init__.py +5 -0
  43. sqlspec/adapters/bigquery/adk/store.py +585 -0
  44. sqlspec/adapters/bigquery/config.py +36 -11
  45. sqlspec/adapters/bigquery/data_dictionary.py +42 -2
  46. sqlspec/adapters/bigquery/driver.py +489 -144
  47. sqlspec/adapters/bigquery/litestar/__init__.py +5 -0
  48. sqlspec/adapters/bigquery/litestar/store.py +327 -0
  49. sqlspec/adapters/bigquery/type_converter.py +55 -23
  50. sqlspec/adapters/duckdb/_types.py +2 -2
  51. sqlspec/adapters/duckdb/adk/__init__.py +14 -0
  52. sqlspec/adapters/duckdb/adk/store.py +563 -0
  53. sqlspec/adapters/duckdb/config.py +79 -21
  54. sqlspec/adapters/duckdb/data_dictionary.py +41 -2
  55. sqlspec/adapters/duckdb/driver.py +225 -44
  56. sqlspec/adapters/duckdb/litestar/__init__.py +5 -0
  57. sqlspec/adapters/duckdb/litestar/store.py +332 -0
  58. sqlspec/adapters/duckdb/pool.py +5 -5
  59. sqlspec/adapters/duckdb/type_converter.py +51 -21
  60. sqlspec/adapters/oracledb/_numpy_handlers.py +133 -0
  61. sqlspec/adapters/oracledb/_types.py +20 -2
  62. sqlspec/adapters/oracledb/adk/__init__.py +5 -0
  63. sqlspec/adapters/oracledb/adk/store.py +1628 -0
  64. sqlspec/adapters/oracledb/config.py +120 -36
  65. sqlspec/adapters/oracledb/data_dictionary.py +87 -20
  66. sqlspec/adapters/oracledb/driver.py +475 -86
  67. sqlspec/adapters/oracledb/litestar/__init__.py +5 -0
  68. sqlspec/adapters/oracledb/litestar/store.py +765 -0
  69. sqlspec/adapters/oracledb/migrations.py +316 -25
  70. sqlspec/adapters/oracledb/type_converter.py +91 -16
  71. sqlspec/adapters/psqlpy/_type_handlers.py +44 -0
  72. sqlspec/adapters/psqlpy/_types.py +2 -1
  73. sqlspec/adapters/psqlpy/adk/__init__.py +5 -0
  74. sqlspec/adapters/psqlpy/adk/store.py +483 -0
  75. sqlspec/adapters/psqlpy/config.py +45 -19
  76. sqlspec/adapters/psqlpy/data_dictionary.py +48 -2
  77. sqlspec/adapters/psqlpy/driver.py +108 -41
  78. sqlspec/adapters/psqlpy/litestar/__init__.py +5 -0
  79. sqlspec/adapters/psqlpy/litestar/store.py +272 -0
  80. sqlspec/adapters/psqlpy/type_converter.py +40 -11
  81. sqlspec/adapters/psycopg/_type_handlers.py +80 -0
  82. sqlspec/adapters/psycopg/_types.py +2 -1
  83. sqlspec/adapters/psycopg/adk/__init__.py +5 -0
  84. sqlspec/adapters/psycopg/adk/store.py +962 -0
  85. sqlspec/adapters/psycopg/config.py +65 -37
  86. sqlspec/adapters/psycopg/data_dictionary.py +91 -3
  87. sqlspec/adapters/psycopg/driver.py +200 -78
  88. sqlspec/adapters/psycopg/litestar/__init__.py +5 -0
  89. sqlspec/adapters/psycopg/litestar/store.py +554 -0
  90. sqlspec/adapters/sqlite/__init__.py +2 -1
  91. sqlspec/adapters/sqlite/_type_handlers.py +86 -0
  92. sqlspec/adapters/sqlite/_types.py +1 -1
  93. sqlspec/adapters/sqlite/adk/__init__.py +5 -0
  94. sqlspec/adapters/sqlite/adk/store.py +582 -0
  95. sqlspec/adapters/sqlite/config.py +85 -16
  96. sqlspec/adapters/sqlite/data_dictionary.py +34 -2
  97. sqlspec/adapters/sqlite/driver.py +120 -52
  98. sqlspec/adapters/sqlite/litestar/__init__.py +5 -0
  99. sqlspec/adapters/sqlite/litestar/store.py +318 -0
  100. sqlspec/adapters/sqlite/pool.py +5 -5
  101. sqlspec/base.py +45 -26
  102. sqlspec/builder/__init__.py +73 -4
  103. sqlspec/builder/_base.py +91 -58
  104. sqlspec/builder/_column.py +5 -5
  105. sqlspec/builder/_ddl.py +98 -89
  106. sqlspec/builder/_delete.py +5 -4
  107. sqlspec/builder/_dml.py +388 -0
  108. sqlspec/{_sql.py → builder/_factory.py} +41 -44
  109. sqlspec/builder/_insert.py +5 -82
  110. sqlspec/builder/{mixins/_join_operations.py → _join.py} +145 -143
  111. sqlspec/builder/_merge.py +446 -11
  112. sqlspec/builder/_parsing_utils.py +9 -11
  113. sqlspec/builder/_select.py +1313 -25
  114. sqlspec/builder/_update.py +11 -42
  115. sqlspec/cli.py +76 -69
  116. sqlspec/config.py +331 -62
  117. sqlspec/core/__init__.py +5 -4
  118. sqlspec/core/cache.py +18 -18
  119. sqlspec/core/compiler.py +6 -8
  120. sqlspec/core/filters.py +55 -47
  121. sqlspec/core/hashing.py +9 -9
  122. sqlspec/core/parameters.py +76 -45
  123. sqlspec/core/result.py +234 -47
  124. sqlspec/core/splitter.py +16 -17
  125. sqlspec/core/statement.py +32 -31
  126. sqlspec/core/type_conversion.py +3 -2
  127. sqlspec/driver/__init__.py +1 -3
  128. sqlspec/driver/_async.py +183 -160
  129. sqlspec/driver/_common.py +197 -109
  130. sqlspec/driver/_sync.py +189 -161
  131. sqlspec/driver/mixins/_result_tools.py +20 -236
  132. sqlspec/driver/mixins/_sql_translator.py +4 -4
  133. sqlspec/exceptions.py +70 -7
  134. sqlspec/extensions/adk/__init__.py +53 -0
  135. sqlspec/extensions/adk/_types.py +51 -0
  136. sqlspec/extensions/adk/converters.py +172 -0
  137. sqlspec/extensions/adk/migrations/0001_create_adk_tables.py +144 -0
  138. sqlspec/extensions/adk/migrations/__init__.py +0 -0
  139. sqlspec/extensions/adk/service.py +181 -0
  140. sqlspec/extensions/adk/store.py +536 -0
  141. sqlspec/extensions/aiosql/adapter.py +69 -61
  142. sqlspec/extensions/fastapi/__init__.py +21 -0
  143. sqlspec/extensions/fastapi/extension.py +331 -0
  144. sqlspec/extensions/fastapi/providers.py +543 -0
  145. sqlspec/extensions/flask/__init__.py +36 -0
  146. sqlspec/extensions/flask/_state.py +71 -0
  147. sqlspec/extensions/flask/_utils.py +40 -0
  148. sqlspec/extensions/flask/extension.py +389 -0
  149. sqlspec/extensions/litestar/__init__.py +21 -4
  150. sqlspec/extensions/litestar/cli.py +54 -10
  151. sqlspec/extensions/litestar/config.py +56 -266
  152. sqlspec/extensions/litestar/handlers.py +46 -17
  153. sqlspec/extensions/litestar/migrations/0001_create_session_table.py +137 -0
  154. sqlspec/extensions/litestar/migrations/__init__.py +3 -0
  155. sqlspec/extensions/litestar/plugin.py +349 -224
  156. sqlspec/extensions/litestar/providers.py +25 -25
  157. sqlspec/extensions/litestar/store.py +265 -0
  158. sqlspec/extensions/starlette/__init__.py +10 -0
  159. sqlspec/extensions/starlette/_state.py +25 -0
  160. sqlspec/extensions/starlette/_utils.py +52 -0
  161. sqlspec/extensions/starlette/extension.py +254 -0
  162. sqlspec/extensions/starlette/middleware.py +154 -0
  163. sqlspec/loader.py +30 -49
  164. sqlspec/migrations/base.py +200 -76
  165. sqlspec/migrations/commands.py +591 -62
  166. sqlspec/migrations/context.py +6 -9
  167. sqlspec/migrations/fix.py +199 -0
  168. sqlspec/migrations/loaders.py +47 -19
  169. sqlspec/migrations/runner.py +241 -75
  170. sqlspec/migrations/tracker.py +237 -21
  171. sqlspec/migrations/utils.py +51 -3
  172. sqlspec/migrations/validation.py +177 -0
  173. sqlspec/protocols.py +106 -36
  174. sqlspec/storage/_utils.py +85 -0
  175. sqlspec/storage/backends/fsspec.py +133 -107
  176. sqlspec/storage/backends/local.py +78 -51
  177. sqlspec/storage/backends/obstore.py +276 -168
  178. sqlspec/storage/registry.py +75 -39
  179. sqlspec/typing.py +30 -84
  180. sqlspec/utils/__init__.py +25 -4
  181. sqlspec/utils/arrow_helpers.py +81 -0
  182. sqlspec/utils/config_resolver.py +6 -6
  183. sqlspec/utils/correlation.py +4 -5
  184. sqlspec/utils/data_transformation.py +3 -2
  185. sqlspec/utils/deprecation.py +9 -8
  186. sqlspec/utils/fixtures.py +4 -4
  187. sqlspec/utils/logging.py +46 -6
  188. sqlspec/utils/module_loader.py +205 -5
  189. sqlspec/utils/portal.py +311 -0
  190. sqlspec/utils/schema.py +288 -0
  191. sqlspec/utils/serializers.py +113 -4
  192. sqlspec/utils/sync_tools.py +36 -22
  193. sqlspec/utils/text.py +1 -2
  194. sqlspec/utils/type_guards.py +136 -20
  195. sqlspec/utils/version.py +433 -0
  196. {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/METADATA +41 -22
  197. sqlspec-0.28.0.dist-info/RECORD +221 -0
  198. sqlspec/builder/mixins/__init__.py +0 -55
  199. sqlspec/builder/mixins/_cte_and_set_ops.py +0 -253
  200. sqlspec/builder/mixins/_delete_operations.py +0 -50
  201. sqlspec/builder/mixins/_insert_operations.py +0 -282
  202. sqlspec/builder/mixins/_merge_operations.py +0 -698
  203. sqlspec/builder/mixins/_order_limit_operations.py +0 -145
  204. sqlspec/builder/mixins/_pivot_operations.py +0 -157
  205. sqlspec/builder/mixins/_select_operations.py +0 -930
  206. sqlspec/builder/mixins/_update_operations.py +0 -199
  207. sqlspec/builder/mixins/_where_clause.py +0 -1298
  208. sqlspec-0.26.0.dist-info/RECORD +0 -157
  209. sqlspec-0.26.0.dist-info/licenses/NOTICE +0 -29
  210. {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/WHEEL +0 -0
  211. {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/entry_points.txt +0 -0
  212. {sqlspec-0.26.0.dist-info → sqlspec-0.28.0.dist-info}/licenses/LICENSE +0 -0
@@ -1,698 +0,0 @@
1
- # pyright: reportPrivateUsage=false
2
- """MERGE operation mixins.
3
-
4
- Provides mixins for MERGE statement functionality including INTO,
5
- USING, ON, WHEN MATCHED, and WHEN NOT MATCHED clauses.
6
- """
7
-
8
- from typing import Any, Optional, Union
9
-
10
- from mypy_extensions import trait
11
- from sqlglot import exp
12
- from typing_extensions import Self
13
-
14
- from sqlspec.builder._parsing_utils import extract_sql_object_expression
15
- from sqlspec.exceptions import SQLBuilderError
16
- from sqlspec.utils.type_guards import has_query_builder_parameters
17
-
18
- __all__ = (
19
- "MergeIntoClauseMixin",
20
- "MergeMatchedClauseMixin",
21
- "MergeNotMatchedBySourceClauseMixin",
22
- "MergeNotMatchedClauseMixin",
23
- "MergeOnClauseMixin",
24
- "MergeUsingClauseMixin",
25
- )
26
-
27
-
28
- @trait
29
- class MergeIntoClauseMixin:
30
- """Mixin providing INTO clause for MERGE builders."""
31
-
32
- __slots__ = ()
33
-
34
- def get_expression(self) -> Optional[exp.Expression]: ...
35
- def set_expression(self, expression: exp.Expression) -> None: ...
36
-
37
- def into(self, table: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
38
- """Set the target table for the MERGE operation (INTO clause).
39
-
40
- Args:
41
- table: The target table name or expression for the MERGE operation.
42
- Can be a string (table name) or an sqlglot Expression.
43
- alias: Optional alias for the target table.
44
-
45
- Returns:
46
- The current builder instance for method chaining.
47
- """
48
- current_expr = self.get_expression()
49
- if current_expr is None or not isinstance(current_expr, exp.Merge):
50
- self.set_expression(exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[])))
51
- current_expr = self.get_expression()
52
-
53
- assert current_expr is not None
54
- current_expr.set("this", exp.to_table(table, alias=alias) if isinstance(table, str) else table)
55
- return self
56
-
57
-
58
- @trait
59
- class MergeUsingClauseMixin:
60
- """Mixin providing USING clause for MERGE builders."""
61
-
62
- __slots__ = ()
63
-
64
- def get_expression(self) -> Optional[exp.Expression]: ...
65
- def set_expression(self, expression: exp.Expression) -> None: ...
66
-
67
- def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
68
- """Add parameter - provided by QueryBuilder."""
69
- msg = "Method must be provided by QueryBuilder subclass"
70
- raise NotImplementedError(msg)
71
-
72
- def _generate_unique_parameter_name(self, base_name: str) -> str:
73
- """Generate unique parameter name - provided by QueryBuilder."""
74
- msg = "Method must be provided by QueryBuilder subclass"
75
- raise NotImplementedError(msg)
76
-
77
- def using(self, source: Union[str, exp.Expression, Any], alias: Optional[str] = None) -> Self:
78
- """Set the source data for the MERGE operation (USING clause).
79
-
80
- Args:
81
- source: The source data for the MERGE operation.
82
- Can be a string (table name), an sqlglot Expression, or a SelectBuilder instance.
83
- alias: Optional alias for the source table.
84
-
85
- Returns:
86
- The current builder instance for method chaining.
87
-
88
- Raises:
89
- SQLBuilderError: If the current expression is not a MERGE statement or if the source type is unsupported.
90
- """
91
- current_expr = self.get_expression()
92
- if current_expr is None or not isinstance(current_expr, exp.Merge):
93
- self.set_expression(exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[])))
94
- current_expr = self.get_expression()
95
-
96
- assert current_expr is not None
97
- source_expr: exp.Expression
98
- if isinstance(source, str):
99
- source_expr = exp.to_table(source, alias=alias)
100
- elif isinstance(source, dict):
101
- columns = list(source.keys())
102
- values = list(source.values())
103
-
104
- parameterized_values: list[exp.Expression] = []
105
- for col, val in zip(columns, values):
106
- column_name = col if isinstance(col, str) else str(col)
107
- if "." in column_name:
108
- column_name = column_name.split(".")[-1]
109
- param_name = self._generate_unique_parameter_name(column_name)
110
- _, param_name = self.add_parameter(val, name=param_name)
111
- parameterized_values.append(exp.Placeholder(this=param_name))
112
-
113
- select_expr = exp.Select()
114
- select_expressions = []
115
- for i, col in enumerate(columns):
116
- select_expressions.append(exp.alias_(parameterized_values[i], col))
117
- select_expr.set("expressions", select_expressions)
118
-
119
- from_expr = exp.From(this=exp.to_table("DUAL"))
120
- select_expr.set("from", from_expr)
121
-
122
- source_expr = exp.paren(select_expr)
123
- if alias:
124
- source_expr = exp.alias_(source_expr, alias, table=False)
125
- elif has_query_builder_parameters(source) and hasattr(source, "_expression"):
126
- subquery_builder_parameters = source.parameters
127
- if subquery_builder_parameters:
128
- for p_name, p_value in subquery_builder_parameters.items():
129
- self.add_parameter(p_value, name=p_name)
130
-
131
- subquery_exp = exp.paren(getattr(source, "_expression", exp.select()))
132
- source_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
133
- elif isinstance(source, exp.Expression):
134
- source_expr = source
135
- if alias:
136
- source_expr = exp.alias_(source_expr, alias)
137
- else:
138
- msg = f"Unsupported source type for USING clause: {type(source)}"
139
- raise SQLBuilderError(msg)
140
-
141
- current_expr.set("using", source_expr)
142
- return self
143
-
144
-
145
- @trait
146
- class MergeOnClauseMixin:
147
- """Mixin providing ON clause for MERGE builders."""
148
-
149
- __slots__ = ()
150
-
151
- def get_expression(self) -> Optional[exp.Expression]: ...
152
- def set_expression(self, expression: exp.Expression) -> None: ...
153
-
154
- def on(self, condition: Union[str, exp.Expression]) -> Self:
155
- """Set the join condition for the MERGE operation (ON clause).
156
-
157
- Args:
158
- condition: The join condition for the MERGE operation.
159
- Can be a string (SQL condition) or an sqlglot Expression.
160
-
161
- Returns:
162
- The current builder instance for method chaining.
163
-
164
- Raises:
165
- SQLBuilderError: If the current expression is not a MERGE statement or if the condition type is unsupported.
166
- """
167
- current_expr = self.get_expression()
168
- if current_expr is None or not isinstance(current_expr, exp.Merge):
169
- self.set_expression(exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[])))
170
- current_expr = self.get_expression()
171
-
172
- assert current_expr is not None
173
- condition_expr: exp.Expression
174
- if isinstance(condition, str):
175
- parsed_condition: Optional[exp.Expression] = exp.maybe_parse(
176
- condition, dialect=getattr(self, "dialect", None)
177
- )
178
- if not parsed_condition:
179
- msg = f"Could not parse ON condition: {condition}"
180
- raise SQLBuilderError(msg)
181
- condition_expr = parsed_condition
182
- elif isinstance(condition, exp.Expression):
183
- condition_expr = condition
184
- else:
185
- msg = f"Unsupported condition type for ON clause: {type(condition)}"
186
- raise SQLBuilderError(msg)
187
-
188
- current_expr.set("on", condition_expr)
189
- return self
190
-
191
-
192
- @trait
193
- class MergeMatchedClauseMixin:
194
- """Mixin providing WHEN MATCHED THEN ... clauses for MERGE builders."""
195
-
196
- __slots__ = ()
197
-
198
- def get_expression(self) -> Optional[exp.Expression]: ...
199
- def set_expression(self, expression: exp.Expression) -> None: ...
200
-
201
- def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
202
- """Add parameter - provided by QueryBuilder."""
203
- msg = "Method must be provided by QueryBuilder subclass"
204
- raise NotImplementedError(msg)
205
-
206
- def _generate_unique_parameter_name(self, base_name: str) -> str:
207
- """Generate unique parameter name - provided by QueryBuilder."""
208
- msg = "Method must be provided by QueryBuilder subclass"
209
- raise NotImplementedError(msg)
210
-
211
- def _is_column_reference(self, value: str) -> bool:
212
- """Check if a string value is a column reference rather than a literal.
213
-
214
- Uses sqlglot to parse the value and determine if it represents a column
215
- reference, function call, or other SQL expression rather than a literal.
216
- """
217
- if not isinstance(value, str):
218
- return False
219
-
220
- # If the string contains spaces and no SQL-like syntax, treat as literal
221
- if " " in value and not any(x in value for x in [".", "(", ")", "*", "="]):
222
- return False
223
-
224
- # Only consider strings with dots (table.column), functions, or SQL keywords as column references
225
- # Simple identifiers are treated as literals
226
- if not any(x in value for x in [".", "(", ")"]):
227
- # Check if it's a SQL keyword/function that should be treated as expression
228
- sql_keywords = {"NULL", "CURRENT_TIMESTAMP", "CURRENT_DATE", "CURRENT_TIME", "DEFAULT"}
229
- if value.upper() not in sql_keywords:
230
- return False
231
-
232
- try:
233
- # Try to parse as SQL expression
234
- parsed: Optional[exp.Expression] = exp.maybe_parse(value)
235
- if parsed is None:
236
- return False
237
-
238
- # Check for SQL literals that should be treated as expressions
239
- return isinstance(
240
- parsed,
241
- (exp.Dot, exp.Anonymous, exp.Func, exp.Null, exp.CurrentTimestamp, exp.CurrentDate, exp.CurrentTime),
242
- )
243
- except Exception:
244
- # If parsing fails, treat as literal
245
- return False
246
-
247
- def _add_when_clause(self, when_clause: exp.When) -> None:
248
- """Helper to add a WHEN clause to the MERGE statement.
249
-
250
- Args:
251
- when_clause: The WHEN clause to add.
252
- """
253
- current_expr = self.get_expression()
254
- if current_expr is None or not isinstance(current_expr, exp.Merge):
255
- self.set_expression(exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[])))
256
- current_expr = self.get_expression()
257
-
258
- assert current_expr is not None
259
- whens = current_expr.args.get("whens")
260
- if not whens:
261
- whens = exp.Whens(expressions=[])
262
- current_expr.set("whens", whens)
263
-
264
- whens.append("expressions", when_clause)
265
-
266
- def when_matched_then_update(
267
- self,
268
- set_values: Optional[dict[str, Any]] = None,
269
- condition: Optional[Union[str, exp.Expression]] = None,
270
- **kwargs: Any,
271
- ) -> Self:
272
- """Define the UPDATE action for matched rows.
273
-
274
- Supports:
275
- - when_matched_then_update({"column": value})
276
- - when_matched_then_update(column=value, other_column=other_value)
277
- - when_matched_then_update({"column": value}, other_column=other_value)
278
-
279
- Args:
280
- set_values: A dictionary of column names and their new values to set.
281
- The values will be parameterized.
282
- condition: An optional additional condition for this specific action.
283
- **kwargs: Column-value pairs to update on match.
284
-
285
- Raises:
286
- SQLBuilderError: If the condition type is unsupported.
287
-
288
- Returns:
289
- The current builder instance for method chaining.
290
- """
291
- # Combine set_values dict and kwargs
292
- all_values = {}
293
- if set_values:
294
- all_values.update(set_values)
295
- if kwargs:
296
- all_values.update(kwargs)
297
-
298
- if not all_values:
299
- msg = "No update values provided. Use set_values dict or kwargs."
300
- raise SQLBuilderError(msg)
301
-
302
- update_expressions: list[exp.EQ] = []
303
- for col, val in all_values.items():
304
- if hasattr(val, "expression") and hasattr(val, "sql"):
305
- value_expr = extract_sql_object_expression(val, builder=self)
306
- elif isinstance(val, exp.Expression):
307
- value_expr = val
308
- elif isinstance(val, str) and self._is_column_reference(val):
309
- value_expr = exp.maybe_parse(val) or exp.column(val)
310
- else:
311
- column_name = col if isinstance(col, str) else str(col)
312
- if "." in column_name:
313
- column_name = column_name.split(".")[-1]
314
- param_name = self._generate_unique_parameter_name(column_name)
315
- _, param_name = self.add_parameter(val, name=param_name)
316
- value_expr = exp.Placeholder(this=param_name)
317
-
318
- update_expressions.append(exp.EQ(this=exp.column(col), expression=value_expr))
319
-
320
- when_args: dict[str, Any] = {"matched": True, "then": exp.Update(expressions=update_expressions)}
321
-
322
- if condition:
323
- condition_expr: exp.Expression
324
- if isinstance(condition, str):
325
- parsed_cond: Optional[exp.Expression] = exp.maybe_parse(
326
- condition, dialect=getattr(self, "dialect", None)
327
- )
328
- if not parsed_cond:
329
- msg = f"Could not parse WHEN clause condition: {condition}"
330
- raise SQLBuilderError(msg)
331
- condition_expr = parsed_cond
332
- elif isinstance(condition, exp.Expression):
333
- condition_expr = condition
334
- else:
335
- msg = f"Unsupported condition type for WHEN clause: {type(condition)}"
336
- raise SQLBuilderError(msg)
337
- when_args["this"] = condition_expr
338
-
339
- when_clause = exp.When(**when_args)
340
- self._add_when_clause(when_clause)
341
- return self
342
-
343
- def when_matched_then_delete(self, condition: Optional[Union[str, exp.Expression]] = None) -> Self:
344
- """Define the DELETE action for matched rows.
345
-
346
- Args:
347
- condition: An optional additional condition for this specific action.
348
-
349
- Raises:
350
- SQLBuilderError: If the condition type is unsupported.
351
-
352
- Returns:
353
- The current builder instance for method chaining.
354
- """
355
- when_args: dict[str, Any] = {"matched": True, "then": exp.Delete()}
356
-
357
- if condition:
358
- condition_expr: exp.Expression
359
- if isinstance(condition, str):
360
- parsed_cond: Optional[exp.Expression] = exp.maybe_parse(
361
- condition, dialect=getattr(self, "dialect", None)
362
- )
363
- if not parsed_cond:
364
- msg = f"Could not parse WHEN clause condition: {condition}"
365
- raise SQLBuilderError(msg)
366
- condition_expr = parsed_cond
367
- elif isinstance(condition, exp.Expression):
368
- condition_expr = condition
369
- else:
370
- msg = f"Unsupported condition type for WHEN clause: {type(condition)}"
371
- raise SQLBuilderError(msg)
372
- when_args["this"] = condition_expr
373
-
374
- when_clause = exp.When(**when_args)
375
- self._add_when_clause(when_clause)
376
- return self
377
-
378
-
379
- @trait
380
- class MergeNotMatchedClauseMixin:
381
- """Mixin providing WHEN NOT MATCHED THEN ... clauses for MERGE builders."""
382
-
383
- __slots__ = ()
384
-
385
- def get_expression(self) -> Optional[exp.Expression]: ...
386
- def set_expression(self, expression: exp.Expression) -> None: ...
387
-
388
- def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
389
- """Add parameter - provided by QueryBuilder."""
390
- msg = "Method must be provided by QueryBuilder subclass"
391
- raise NotImplementedError(msg)
392
-
393
- def _generate_unique_parameter_name(self, base_name: str) -> str:
394
- """Generate unique parameter name - provided by QueryBuilder."""
395
- msg = "Method must be provided by QueryBuilder subclass"
396
- raise NotImplementedError(msg)
397
-
398
- def _is_column_reference(self, value: str) -> bool:
399
- """Check if a string value is a column reference rather than a literal.
400
-
401
- Uses sqlglot to parse the value and determine if it represents a column
402
- reference, function call, or other SQL expression rather than a literal.
403
- """
404
- if not isinstance(value, str):
405
- return False
406
-
407
- # If the string contains spaces and no SQL-like syntax, treat as literal
408
- if " " in value and not any(x in value for x in [".", "(", ")", "*", "="]):
409
- return False
410
-
411
- try:
412
- # Try to parse as SQL expression
413
- parsed: Optional[exp.Expression] = exp.maybe_parse(value)
414
- if parsed is None:
415
- return False
416
-
417
- # If it parses to a Column, Dot (table.column), Identifier, or other SQL constructs
418
-
419
- except Exception:
420
- # If parsing fails, fall back to conservative approach
421
- # Only treat simple identifiers as column references
422
- return (
423
- value.replace("_", "").replace(".", "").isalnum()
424
- and (value[0].isalpha() or value[0] == "_")
425
- and " " not in value
426
- and "'" not in value
427
- and '"' not in value
428
- )
429
- return bool(
430
- isinstance(
431
- parsed,
432
- (
433
- exp.Column,
434
- exp.Dot,
435
- exp.Identifier,
436
- exp.Anonymous,
437
- exp.Func,
438
- exp.Null,
439
- exp.CurrentTimestamp,
440
- exp.CurrentDate,
441
- exp.CurrentTime,
442
- ),
443
- )
444
- )
445
-
446
- def _add_when_clause(self, when_clause: exp.When) -> None:
447
- """Helper to add a WHEN clause to the MERGE statement - provided by QueryBuilder."""
448
- msg = "Method must be provided by QueryBuilder subclass"
449
- raise NotImplementedError(msg)
450
-
451
- def when_not_matched_then_insert(
452
- self,
453
- columns: Optional[list[str]] = None,
454
- values: Optional[list[Any]] = None,
455
- condition: Optional[Union[str, exp.Expression]] = None,
456
- by_target: bool = True,
457
- ) -> Self:
458
- """Define the INSERT action for rows not matched.
459
-
460
- Args:
461
- columns: A list of column names to insert into. If None, implies INSERT DEFAULT VALUES or matching source columns.
462
- values: A list of values corresponding to the columns.
463
- These values will be parameterized. If None, implies INSERT DEFAULT VALUES or subquery source.
464
- condition: An optional additional condition for this specific action.
465
- by_target: If True (default), condition is "WHEN NOT MATCHED [BY TARGET]".
466
- If False, condition is "WHEN NOT MATCHED BY SOURCE".
467
-
468
- Returns:
469
- The current builder instance for method chaining.
470
-
471
- Raises:
472
- SQLBuilderError: If columns and values are provided but do not match in length,
473
- or if columns are provided without values.
474
- """
475
- insert_args: dict[str, Any] = {}
476
- if columns and values:
477
- if len(columns) != len(values):
478
- msg = "Number of columns must match number of values for INSERT."
479
- raise SQLBuilderError(msg)
480
-
481
- parameterized_values: list[exp.Expression] = []
482
- for i, val in enumerate(values):
483
- if isinstance(val, str) and self._is_column_reference(val):
484
- # Handle column references (like "s.data") as column expressions, not parameters
485
- parameterized_values.append(exp.maybe_parse(val) or exp.column(val))
486
- else:
487
- column_name = columns[i] if isinstance(columns[i], str) else str(columns[i])
488
- if "." in column_name:
489
- column_name = column_name.split(".")[-1]
490
- param_name = self._generate_unique_parameter_name(column_name)
491
- _, param_name = self.add_parameter(val, name=param_name)
492
- parameterized_values.append(exp.Placeholder(this=param_name))
493
-
494
- insert_args["this"] = exp.Tuple(expressions=[exp.column(c) for c in columns])
495
- insert_args["expression"] = exp.Tuple(expressions=parameterized_values)
496
- elif columns and not values:
497
- msg = "Specifying columns without values for INSERT action is complex and not fully supported yet. Consider providing full expressions."
498
- raise SQLBuilderError(msg)
499
- elif not columns and not values:
500
- pass
501
- else:
502
- msg = "Cannot specify values without columns for INSERT action."
503
- raise SQLBuilderError(msg)
504
-
505
- when_args: dict[str, Any] = {"matched": False, "then": exp.Insert(**insert_args)}
506
-
507
- if not by_target:
508
- when_args["source"] = True
509
-
510
- if condition:
511
- condition_expr: exp.Expression
512
- if isinstance(condition, str):
513
- parsed_cond: Optional[exp.Expression] = exp.maybe_parse(
514
- condition, dialect=getattr(self, "dialect", None)
515
- )
516
- if not parsed_cond:
517
- msg = f"Could not parse WHEN clause condition: {condition}"
518
- raise SQLBuilderError(msg)
519
- condition_expr = parsed_cond
520
- elif isinstance(condition, exp.Expression):
521
- condition_expr = condition
522
- else:
523
- msg = f"Unsupported condition type for WHEN clause: {type(condition)}"
524
- raise SQLBuilderError(msg)
525
- when_args["this"] = condition_expr
526
-
527
- when_clause = exp.When(**when_args)
528
- self._add_when_clause(when_clause)
529
- return self
530
-
531
-
532
- @trait
533
- class MergeNotMatchedBySourceClauseMixin:
534
- """Mixin providing WHEN NOT MATCHED BY SOURCE THEN ... clauses for MERGE builders."""
535
-
536
- __slots__ = ()
537
-
538
- def get_expression(self) -> Optional[exp.Expression]: ...
539
- def set_expression(self, expression: exp.Expression) -> None: ...
540
-
541
- def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
542
- """Add parameter - provided by QueryBuilder."""
543
- msg = "Method must be provided by QueryBuilder subclass"
544
- raise NotImplementedError(msg)
545
-
546
- def _generate_unique_parameter_name(self, base_name: str) -> str:
547
- """Generate unique parameter name - provided by QueryBuilder."""
548
- msg = "Method must be provided by QueryBuilder subclass"
549
- raise NotImplementedError(msg)
550
-
551
- def _add_when_clause(self, when_clause: exp.When) -> None:
552
- """Helper to add a WHEN clause to the MERGE statement - provided by QueryBuilder."""
553
- msg = "Method must be provided by QueryBuilder subclass"
554
- raise NotImplementedError(msg)
555
-
556
- def _is_column_reference(self, value: str) -> bool:
557
- """Check if a string value is a column reference rather than a literal.
558
-
559
- Uses sqlglot to parse the value and determine if it represents a column
560
- reference, function call, or other SQL expression rather than a literal.
561
-
562
- Args:
563
- value: The string value to check
564
-
565
- Returns:
566
- True if the value is a column reference, False if it's a literal
567
- """
568
- if not isinstance(value, str):
569
- return False
570
-
571
- try:
572
- parsed: Optional[exp.Expression] = exp.maybe_parse(value)
573
- if parsed is None:
574
- return False
575
-
576
- except Exception:
577
- return False
578
- return bool(
579
- isinstance(
580
- parsed,
581
- (exp.Dot, exp.Anonymous, exp.Func, exp.Null, exp.CurrentTimestamp, exp.CurrentDate, exp.CurrentTime),
582
- )
583
- )
584
-
585
- def when_not_matched_by_source_then_update(
586
- self,
587
- set_values: Optional[dict[str, Any]] = None,
588
- condition: Optional[Union[str, exp.Expression]] = None,
589
- **kwargs: Any,
590
- ) -> Self:
591
- """Define the UPDATE action for rows not matched by source.
592
-
593
- This is useful for handling rows that exist in the target but not in the source.
594
-
595
- Supports:
596
- - when_not_matched_by_source_then_update({"column": value})
597
- - when_not_matched_by_source_then_update(column=value, other_column=other_value)
598
- - when_not_matched_by_source_then_update({"column": value}, other_column=other_value)
599
-
600
- Args:
601
- set_values: A dictionary of column names and their new values to set.
602
- condition: An optional additional condition for this specific action.
603
- **kwargs: Column-value pairs to update when not matched by source.
604
-
605
- Raises:
606
- SQLBuilderError: If the condition type is unsupported.
607
-
608
- Returns:
609
- The current builder instance for method chaining.
610
- """
611
- # Combine set_values dict and kwargs
612
- all_values = dict(set_values or {}, **kwargs)
613
-
614
- if not all_values:
615
- msg = "No update values provided. Use set_values dict or kwargs."
616
- raise SQLBuilderError(msg)
617
-
618
- update_expressions: list[exp.EQ] = []
619
- for col, val in all_values.items():
620
- if hasattr(val, "expression") and hasattr(val, "sql"):
621
- value_expr = extract_sql_object_expression(val, builder=self)
622
- elif isinstance(val, exp.Expression):
623
- value_expr = val
624
- elif isinstance(val, str) and self._is_column_reference(val):
625
- value_expr = exp.maybe_parse(val) or exp.column(val)
626
- else:
627
- column_name = col if isinstance(col, str) else str(col)
628
- if "." in column_name:
629
- column_name = column_name.split(".")[-1]
630
- param_name = self._generate_unique_parameter_name(column_name)
631
- _, param_name = self.add_parameter(val, name=param_name)
632
- value_expr = exp.Placeholder(this=param_name)
633
-
634
- update_expressions.append(exp.EQ(this=exp.column(col), expression=value_expr))
635
-
636
- when_args: dict[str, Any] = {
637
- "matched": False,
638
- "source": True,
639
- "then": exp.Update(expressions=update_expressions),
640
- }
641
-
642
- if condition:
643
- condition_expr: exp.Expression
644
- if isinstance(condition, str):
645
- parsed_cond: Optional[exp.Expression] = exp.maybe_parse(
646
- condition, dialect=getattr(self, "dialect", None)
647
- )
648
- if not parsed_cond:
649
- msg = f"Could not parse WHEN clause condition: {condition}"
650
- raise SQLBuilderError(msg)
651
- condition_expr = parsed_cond
652
- elif isinstance(condition, exp.Expression):
653
- condition_expr = condition
654
- else:
655
- msg = f"Unsupported condition type for WHEN clause: {type(condition)}"
656
- raise SQLBuilderError(msg)
657
- when_args["this"] = condition_expr
658
-
659
- when_clause = exp.When(**when_args)
660
- self._add_when_clause(when_clause)
661
- return self
662
-
663
- def when_not_matched_by_source_then_delete(self, condition: Optional[Union[str, exp.Expression]] = None) -> Self:
664
- """Define the DELETE action for rows not matched by source.
665
-
666
- This is useful for cleaning up rows that exist in the target but not in the source.
667
-
668
- Args:
669
- condition: An optional additional condition for this specific action.
670
-
671
- Raises:
672
- SQLBuilderError: If the condition type is unsupported.
673
-
674
- Returns:
675
- The current builder instance for method chaining.
676
- """
677
- when_args: dict[str, Any] = {"matched": False, "source": True, "then": exp.Delete()}
678
-
679
- if condition:
680
- condition_expr: exp.Expression
681
- if isinstance(condition, str):
682
- parsed_cond: Optional[exp.Expression] = exp.maybe_parse(
683
- condition, dialect=getattr(self, "dialect", None)
684
- )
685
- if not parsed_cond:
686
- msg = f"Could not parse WHEN clause condition: {condition}"
687
- raise SQLBuilderError(msg)
688
- condition_expr = parsed_cond
689
- elif isinstance(condition, exp.Expression):
690
- condition_expr = condition
691
- else:
692
- msg = f"Unsupported condition type for WHEN clause: {type(condition)}"
693
- raise SQLBuilderError(msg)
694
- when_args["this"] = condition_expr
695
-
696
- when_clause = exp.When(**when_args)
697
- self._add_when_clause(when_clause)
698
- return self