sqlalchemy-cratedb 0.41.0.dev0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,337 @@
1
+ # -*- coding: utf-8; -*-
2
+ #
3
+ # Licensed to CRATE Technology GmbH ("Crate") under one or more contributor
4
+ # license agreements. See the NOTICE file distributed with this work for
5
+ # additional information regarding copyright ownership. Crate licenses
6
+ # this file to you under the Apache License, Version 2.0 (the "License");
7
+ # you may not use this file except in compliance with the License. You may
8
+ # obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing, software
13
+ # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
14
+ # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
15
+ # License for the specific language governing permissions and limitations
16
+ # under the License.
17
+ #
18
+ # However, if you have executed another commercial license agreement
19
+ # with Crate these terms will supersede the license and you may use the
20
+ # software solely pursuant to the terms of the relevant commercial agreement.
21
+
22
+ import sqlalchemy as sa
23
+ from sqlalchemy.dialects.postgresql.base import PGCompiler
24
+ from sqlalchemy.sql import selectable
25
+ from sqlalchemy.sql.crud import (
26
+ REQUIRED,
27
+ _create_bind_param,
28
+ _extend_values_for_multiparams,
29
+ _get_stmt_parameter_tuples_params,
30
+ _get_update_multitable_params,
31
+ _key_getters_for_crud_column,
32
+ _scan_cols,
33
+ _scan_insert_from_select_cols,
34
+ )
35
+
36
+ from sqlalchemy_cratedb.compiler import CrateCompiler
37
+
38
+
39
+ class CrateCompilerSA14(CrateCompiler):
40
+ def returning_clause(self, stmt, returning_cols):
41
+ """
42
+ Generate RETURNING clause, PostgreSQL-compatible.
43
+ """
44
+ return PGCompiler.returning_clause(self, stmt, returning_cols)
45
+
46
+ def visit_update(self, update_stmt, **kw):
47
+ compile_state = update_stmt._compile_state_factory(update_stmt, self, **kw)
48
+ update_stmt = compile_state.statement
49
+
50
+ # [14] CrateDB patch.
51
+ if not compile_state._dict_parameters and not hasattr(update_stmt, "_crate_specific"):
52
+ return super().visit_update(update_stmt, **kw)
53
+
54
+ toplevel = not self.stack
55
+ if toplevel:
56
+ self.isupdate = True
57
+ if not self.compile_state:
58
+ self.compile_state = compile_state
59
+
60
+ extra_froms = compile_state._extra_froms
61
+ is_multitable = bool(extra_froms)
62
+
63
+ if is_multitable:
64
+ # main table might be a JOIN
65
+ main_froms = set(selectable._from_objects(update_stmt.table))
66
+ render_extra_froms = [f for f in extra_froms if f not in main_froms]
67
+ correlate_froms = main_froms.union(extra_froms)
68
+ else:
69
+ render_extra_froms = []
70
+ correlate_froms = {update_stmt.table}
71
+
72
+ self.stack.append(
73
+ {
74
+ "correlate_froms": correlate_froms,
75
+ "asfrom_froms": correlate_froms,
76
+ "selectable": update_stmt,
77
+ }
78
+ )
79
+
80
+ text = "UPDATE "
81
+
82
+ if update_stmt._prefixes:
83
+ text += self._generate_prefixes(update_stmt, update_stmt._prefixes, **kw)
84
+
85
+ table_text = self.update_tables_clause(
86
+ update_stmt, update_stmt.table, render_extra_froms, **kw
87
+ )
88
+
89
+ # [14] CrateDB patch.
90
+ crud_params = _get_crud_params(self, update_stmt, compile_state, **kw)
91
+
92
+ if update_stmt._hints:
93
+ dialect_hints, table_text = self._setup_crud_hints(update_stmt, table_text)
94
+ else:
95
+ dialect_hints = None
96
+
97
+ if update_stmt._independent_ctes:
98
+ for cte in update_stmt._independent_ctes:
99
+ cte._compiler_dispatch(self, **kw)
100
+
101
+ text += table_text
102
+
103
+ text += " SET "
104
+
105
+ # [14] CrateDB patch begin.
106
+ include_table = extra_froms and self.render_table_with_column_in_update_from
107
+
108
+ set_clauses = []
109
+
110
+ for c, expr, value in crud_params: # noqa: B007
111
+ key = c._compiler_dispatch(self, include_table=include_table)
112
+ clause = key + " = " + value
113
+ set_clauses.append(clause)
114
+
115
+ for k, v in compile_state._dict_parameters.items():
116
+ if isinstance(k, str) and "[" in k:
117
+ bindparam = sa.sql.bindparam(k, v)
118
+ clause = k + " = " + self.process(bindparam)
119
+ set_clauses.append(clause)
120
+
121
+ text += ", ".join(set_clauses)
122
+ # [14] CrateDB patch end.
123
+
124
+ if self.returning or update_stmt._returning:
125
+ if self.returning_precedes_values:
126
+ text += " " + self.returning_clause(
127
+ update_stmt, self.returning or update_stmt._returning
128
+ )
129
+
130
+ if extra_froms:
131
+ extra_from_text = self.update_from_clause(
132
+ update_stmt, update_stmt.table, render_extra_froms, dialect_hints, **kw
133
+ )
134
+ if extra_from_text:
135
+ text += " " + extra_from_text
136
+
137
+ if update_stmt._where_criteria:
138
+ t = self._generate_delimited_and_list(update_stmt._where_criteria, **kw)
139
+ if t:
140
+ text += " WHERE " + t
141
+
142
+ limit_clause = self.update_limit_clause(update_stmt)
143
+ if limit_clause:
144
+ text += " " + limit_clause
145
+
146
+ if (self.returning or update_stmt._returning) and not self.returning_precedes_values:
147
+ text += " " + self.returning_clause(
148
+ update_stmt, self.returning or update_stmt._returning
149
+ )
150
+
151
+ if self.ctes:
152
+ nesting_level = len(self.stack) if not toplevel else None
153
+ text = self._render_cte_clause(nesting_level=nesting_level) + text
154
+
155
+ self.stack.pop(-1)
156
+
157
+ return text
158
+
159
+
160
+ def _get_crud_params(compiler, stmt, compile_state, **kw):
161
+ """create a set of tuples representing column/string pairs for use
162
+ in an INSERT or UPDATE statement.
163
+
164
+ Also generates the Compiled object's postfetch, prefetch, and
165
+ returning column collections, used for default handling and ultimately
166
+ populating the CursorResult's prefetch_cols() and postfetch_cols()
167
+ collections.
168
+
169
+ """
170
+
171
+ compiler.postfetch = []
172
+ compiler.insert_prefetch = []
173
+ compiler.update_prefetch = []
174
+ compiler.returning = []
175
+
176
+ # getters - these are normally just column.key,
177
+ # but in the case of mysql multi-table update, the rules for
178
+ # .key must conditionally take tablename into account
179
+ (
180
+ _column_as_key,
181
+ _getattr_col_key,
182
+ _col_bind_name,
183
+ ) = getters = _key_getters_for_crud_column(compiler, stmt, compile_state)
184
+
185
+ compiler._key_getters_for_crud_column = getters
186
+
187
+ # no parameters in the statement, no parameters in the
188
+ # compiled params - return binds for all columns
189
+ if compiler.column_keys is None and compile_state._no_parameters:
190
+ return [
191
+ (
192
+ c,
193
+ compiler.preparer.format_column(c),
194
+ _create_bind_param(compiler, c, None, required=True),
195
+ )
196
+ for c in stmt.table.columns
197
+ ]
198
+
199
+ if compile_state._has_multi_parameters:
200
+ spd = compile_state._multi_parameters[0]
201
+ stmt_parameter_tuples = list(spd.items())
202
+ elif compile_state._ordered_values:
203
+ spd = compile_state._dict_parameters
204
+ stmt_parameter_tuples = compile_state._ordered_values
205
+ elif compile_state._dict_parameters:
206
+ spd = compile_state._dict_parameters
207
+ stmt_parameter_tuples = list(spd.items())
208
+ else:
209
+ stmt_parameter_tuples = spd = None
210
+
211
+ # if we have statement parameters - set defaults in the
212
+ # compiled params
213
+ if compiler.column_keys is None:
214
+ parameters = {}
215
+ elif stmt_parameter_tuples:
216
+ parameters = dict(
217
+ (_column_as_key(key), REQUIRED) for key in compiler.column_keys if key not in spd
218
+ )
219
+ else:
220
+ parameters = dict((_column_as_key(key), REQUIRED) for key in compiler.column_keys)
221
+
222
+ # create a list of column assignment clauses as tuples
223
+ values = []
224
+
225
+ if stmt_parameter_tuples is not None:
226
+ _get_stmt_parameter_tuples_params(
227
+ compiler,
228
+ compile_state,
229
+ parameters,
230
+ stmt_parameter_tuples,
231
+ _column_as_key,
232
+ values,
233
+ kw,
234
+ )
235
+
236
+ check_columns = {}
237
+
238
+ # special logic that only occurs for multi-table UPDATE
239
+ # statements
240
+ if compile_state.isupdate and compile_state.is_multitable:
241
+ _get_update_multitable_params(
242
+ compiler,
243
+ stmt,
244
+ compile_state,
245
+ stmt_parameter_tuples,
246
+ check_columns,
247
+ _col_bind_name,
248
+ _getattr_col_key,
249
+ values,
250
+ kw,
251
+ )
252
+
253
+ if compile_state.isinsert and stmt._select_names:
254
+ _scan_insert_from_select_cols(
255
+ compiler,
256
+ stmt,
257
+ compile_state,
258
+ parameters,
259
+ _getattr_col_key,
260
+ _column_as_key,
261
+ _col_bind_name,
262
+ check_columns,
263
+ values,
264
+ kw,
265
+ )
266
+ else:
267
+ _scan_cols(
268
+ compiler,
269
+ stmt,
270
+ compile_state,
271
+ parameters,
272
+ _getattr_col_key,
273
+ _column_as_key,
274
+ _col_bind_name,
275
+ check_columns,
276
+ values,
277
+ kw,
278
+ )
279
+
280
+ # [14] CrateDB patch.
281
+ #
282
+ # This sanity check performed by SQLAlchemy currently needs to be
283
+ # deactivated in order to satisfy the rewriting logic of the CrateDB
284
+ # dialect in `rewrite_update` and `visit_update`.
285
+ #
286
+ # It can be quickly reproduced by activating this section and running the
287
+ # test cases::
288
+ #
289
+ # ./bin/test -vvvv -t dict_test
290
+ #
291
+ # That croaks like::
292
+ #
293
+ # sqlalchemy.exc.CompileError: Unconsumed column names: characters_name, data['nested']
294
+ #
295
+ # TODO: Investigate why this is actually happening and eventually mitigate
296
+ # the root cause.
297
+ """
298
+ if parameters and stmt_parameter_tuples:
299
+ check = (
300
+ set(parameters)
301
+ .intersection(_column_as_key(k) for k, v in stmt_parameter_tuples)
302
+ .difference(check_columns)
303
+ )
304
+ if check:
305
+ raise exc.CompileError(
306
+ "Unconsumed column names: %s"
307
+ % (", ".join("%s" % (c,) for c in check))
308
+ )
309
+ """
310
+
311
+ if compile_state._has_multi_parameters:
312
+ values = _extend_values_for_multiparams(
313
+ compiler,
314
+ stmt,
315
+ compile_state,
316
+ values,
317
+ _column_as_key,
318
+ kw,
319
+ )
320
+ elif (
321
+ not values
322
+ and compiler.for_executemany # noqa: W503
323
+ and compiler.dialect.supports_default_metavalue # noqa: W503
324
+ ):
325
+ # convert an "INSERT DEFAULT VALUES"
326
+ # into INSERT (firstcol) VALUES (DEFAULT) which can be turned
327
+ # into an in-place multi values. This supports
328
+ # insert_executemany_returning mode :)
329
+ values = [
330
+ (
331
+ stmt.table.columns[0],
332
+ compiler.preparer.format_column(stmt.table.columns[0]),
333
+ "DEFAULT",
334
+ )
335
+ ]
336
+
337
+ return values