1"""sqlglot expressions DML."""
  2
  3from __future__ import annotations
  4
  5import typing as t
  6
  7from sqlglot.helper import trait
  8from sqlglot.expressions.core import (
  9    Expr,
 10    Expression,
 11    _apply_builder,
 12    _apply_list_builder,
 13    maybe_copy,
 14    _apply_conjunction_builder,
 15)
 16from sqlglot.expressions.ddl import DDL
 17from sqlglot.expressions.query import (
 18    Table,
 19    Where,
 20    From,
 21    _apply_cte_builder,
 22)
 23
 24if t.TYPE_CHECKING:
 25    from typing_extensions import Self, Unpack
 26    from sqlglot.dialects.dialect import DialectType
 27    from sqlglot.expressions.core import ExpOrStr
 28    from sqlglot._typing import ParserNoDialectArgs
 29
 30
 31@trait
 32class DML(Expr):
 33    """Trait for data manipulation language statements."""
 34
 35    def returning(
 36        self,
 37        expression: ExpOrStr,
 38        dialect: DialectType = None,
 39        copy: bool = True,
 40        **opts: Unpack[ParserNoDialectArgs],
 41    ) -> Self:
 42        """
 43        Set the RETURNING expression. Not supported by all dialects.
 44
 45        Example:
 46            >>> Delete().delete("tbl").returning("*", dialect="postgres").sql()
 47            'DELETE FROM tbl RETURNING *'
 48
 49        Args:
 50            expression: the SQL code strings to parse.
 51                If an `Expr` instance is passed, it will be used as-is.
 52            dialect: the dialect used to parse the input expressions.
 53            copy: if `False`, modify this expression instance in-place.
 54            opts: other options to use to parse the input expressions.
 55
 56        Returns:
 57            Delete: the modified expression.
 58        """
 59        return _apply_builder(
 60            expression=expression,
 61            instance=self,
 62            arg="returning",
 63            prefix="RETURNING",
 64            dialect=dialect,
 65            copy=copy,
 66            into=Returning,
 67            **opts,
 68        )
 69
 70
 71class Delete(Expression, DML):
 72    arg_types = {
 73        "with_": False,
 74        "this": False,
 75        "using": False,
 76        "where": False,
 77        "returning": False,
 78        "order": False,
 79        "limit": False,
 80        "tables": False,  # Multiple-Table Syntax (MySQL)
 81        "cluster": False,  # Clickhouse
 82        "hint": False,
 83    }
 84
 85    def delete(
 86        self,
 87        table: ExpOrStr,
 88        dialect: DialectType = None,
 89        copy: bool = True,
 90        **opts: Unpack[ParserNoDialectArgs],
 91    ) -> Delete:
 92        """
 93        Create a DELETE expression or replace the table on an existing DELETE expression.
 94
 95        Example:
 96            >>> Delete().delete("tbl").sql()
 97            'DELETE FROM tbl'
 98
 99        Args:
100            table: the table from which to delete.
101            dialect: the dialect used to parse the input expression.
102            copy: if `False`, modify this expression instance in-place.
103            opts: other options to use to parse the input expressions.
104
105        Returns:
106            Delete: the modified expression.
107        """
108        return _apply_builder(
109            expression=table,
110            instance=self,
111            arg="this",
112            dialect=dialect,
113            into=Table,
114            copy=copy,
115            **opts,
116        )
117
118    def where(
119        self,
120        *expressions: ExpOrStr | None,
121        append: bool = True,
122        dialect: DialectType = None,
123        copy: bool = True,
124        **opts: Unpack[ParserNoDialectArgs],
125    ) -> Delete:
126        """
127        Append to or set the WHERE expressions.
128
129        Example:
130            >>> Delete().delete("tbl").where("x = 'a' OR x < 'b'").sql()
131            "DELETE FROM tbl WHERE x = 'a' OR x < 'b'"
132
133        Args:
134            *expressions: the SQL code strings to parse.
135                If an `Expr` instance is passed, it will be used as-is.
136                Multiple expressions are combined with an AND operator.
137            append: if `True`, AND the new expressions to any existing expression.
138                Otherwise, this resets the expression.
139            dialect: the dialect used to parse the input expressions.
140            copy: if `False`, modify this expression instance in-place.
141            opts: other options to use to parse the input expressions.
142
143        Returns:
144            Delete: the modified expression.
145        """
146        return _apply_conjunction_builder(
147            *expressions,
148            instance=self,
149            arg="where",
150            append=append,
151            into=Where,
152            dialect=dialect,
153            copy=copy,
154            **opts,
155        )
156
157
158class Export(Expression):
159    arg_types = {"this": True, "connection": False, "options": True}
160
161
162class CopyParameter(Expression):
163    arg_types = {"this": True, "expression": False, "expressions": False}
164
165
166class Copy(Expression, DML):
167    arg_types = {
168        "this": True,
169        "kind": True,
170        "files": False,
171        "credentials": False,
172        "format": False,
173        "params": False,
174    }
175
176
177class Credentials(Expression):
178    arg_types = {
179        "credentials": False,
180        "encryption": False,
181        "storage": False,
182        "iam_role": False,
183        "region": False,
184    }
185
186
187class Directory(Expression):
188    arg_types = {"this": True, "local": False, "row_format": False}
189
190
191class DirectoryStage(Expression):
192    pass
193
194
195class Insert(Expression, DDL, DML):
196    arg_types = {
197        "hint": False,
198        "with_": False,
199        "is_function": False,
200        "this": False,
201        "expression": False,
202        "conflict": False,
203        "returning": False,
204        "overwrite": False,
205        "exists": False,
206        "alternative": False,
207        "where": False,
208        "ignore": False,
209        "by_name": False,
210        "stored": False,
211        "partition": False,
212        "settings": False,
213        "source": False,
214        "default": False,
215        "using": False,
216    }
217
218    def with_(
219        self,
220        alias: ExpOrStr,
221        as_: ExpOrStr,
222        recursive: bool | None = None,
223        materialized: bool | None = None,
224        append: bool = True,
225        dialect: DialectType = None,
226        copy: bool = True,
227        **opts: Unpack[ParserNoDialectArgs],
228    ) -> Insert:
229        """
230        Append to or set the common table expressions.
231
232        Example:
233            >>> import sqlglot
234            >>> sqlglot.parse_one("INSERT INTO t SELECT x FROM cte").with_("cte", as_="SELECT * FROM tbl").sql()
235            'WITH cte AS (SELECT * FROM tbl) INSERT INTO t SELECT x FROM cte'
236
237        Args:
238            alias: the SQL code string to parse as the table name.
239                If an `Expr` instance is passed, this is used as-is.
240            as_: the SQL code string to parse as the table expression.
241                If an `Expr` instance is passed, it will be used as-is.
242            recursive: set the RECURSIVE part of the expression. Defaults to `False`.
243            materialized: set the MATERIALIZED part of the expression.
244            append: if `True`, add to any existing expressions.
245                Otherwise, this resets the expressions.
246            dialect: the dialect used to parse the input expression.
247            copy: if `False`, modify this expression instance in-place.
248            opts: other options to use to parse the input expressions.
249
250        Returns:
251            The modified expression.
252        """
253        return _apply_cte_builder(
254            self,
255            alias,
256            as_,
257            recursive=recursive,
258            materialized=materialized,
259            append=append,
260            dialect=dialect,
261            copy=copy,
262            **opts,
263        )
264
265
266class OnConflict(Expression):
267    arg_types = {
268        "duplicate": False,
269        "expressions": False,
270        "action": False,
271        "conflict_keys": False,
272        "index_predicate": False,
273        "constraint": False,
274        "where": False,
275    }
276
277
278class Returning(Expression):
279    arg_types = {"expressions": True, "into": False}
280
281
282class LoadData(Expression):
283    arg_types = {
284        "this": True,
285        "local": False,
286        "overwrite": False,
287        "temp": False,
288        "inpath": False,
289        "files": False,
290        "partition": False,
291        "input_format": False,
292        "serde": False,
293    }
294
295
296class Update(Expression, DML):
297    arg_types = {
298        "with_": False,
299        "this": False,
300        "expressions": False,
301        "from_": False,
302        "where": False,
303        "returning": False,
304        "order": False,
305        "limit": False,
306        "options": False,
307        "hint": False,
308    }
309
310    def table(
311        self,
312        expression: ExpOrStr,
313        dialect: DialectType = None,
314        copy: bool = True,
315        **opts: Unpack[ParserNoDialectArgs],
316    ) -> Update:
317        """
318        Set the table to update.
319
320        Example:
321            >>> Update().table("my_table").set_("x = 1").sql()
322            'UPDATE my_table SET x = 1'
323
324        Args:
325            expression : the SQL code strings to parse.
326                If a `Table` instance is passed, this is used as-is.
327                If another `Expr` instance is passed, it will be wrapped in a `Table`.
328            dialect: the dialect used to parse the input expression.
329            copy: if `False`, modify this expression instance in-place.
330            opts: other options to use to parse the input expressions.
331
332        Returns:
333            The modified Update expression.
334        """
335        return _apply_builder(
336            expression=expression,
337            instance=self,
338            arg="this",
339            into=Table,
340            prefix=None,
341            dialect=dialect,
342            copy=copy,
343            **opts,
344        )
345
346    def set_(
347        self,
348        *expressions: ExpOrStr,
349        append: bool = True,
350        dialect: DialectType = None,
351        copy: bool = True,
352        **opts: Unpack[ParserNoDialectArgs],
353    ) -> Update:
354        """
355        Append to or set the SET expressions.
356
357        Example:
358            >>> Update().table("my_table").set_("x = 1").sql()
359            'UPDATE my_table SET x = 1'
360
361        Args:
362            *expressions: the SQL code strings to parse.
363                If `Expr` instance(s) are passed, they will be used as-is.
364                Multiple expressions are combined with a comma.
365            append: if `True`, add the new expressions to any existing SET expressions.
366                Otherwise, this resets the expressions.
367            dialect: the dialect used to parse the input expressions.
368            copy: if `False`, modify this expression instance in-place.
369            opts: other options to use to parse the input expressions.
370        """
371        return _apply_list_builder(
372            *expressions,
373            instance=self,
374            arg="expressions",
375            append=append,
376            into=Expr,
377            prefix=None,
378            dialect=dialect,
379            copy=copy,
380            **opts,
381        )
382
383    def where(
384        self,
385        *expressions: ExpOrStr | None,
386        append: bool = True,
387        dialect: DialectType = None,
388        copy: bool = True,
389        **opts: Unpack[ParserNoDialectArgs],
390    ) -> Update:
391        """
392        Append to or set the WHERE expressions.
393
394        Example:
395            >>> Update().table("tbl").set_("x = 1").where("x = 'a' OR x < 'b'").sql()
396            "UPDATE tbl SET x = 1 WHERE x = 'a' OR x < 'b'"
397
398        Args:
399            *expressions: the SQL code strings to parse.
400                If an `Expr` instance is passed, it will be used as-is.
401                Multiple expressions are combined with an AND operator.
402            append: if `True`, AND the new expressions to any existing expression.
403                Otherwise, this resets the expression.
404            dialect: the dialect used to parse the input expressions.
405            copy: if `False`, modify this expression instance in-place.
406            opts: other options to use to parse the input expressions.
407
408        Returns:
409            Update: the modified expression.
410        """
411        return _apply_conjunction_builder(
412            *expressions,
413            instance=self,
414            arg="where",
415            append=append,
416            into=Where,
417            dialect=dialect,
418            copy=copy,
419            **opts,
420        )
421
422    def from_(
423        self,
424        expression: ExpOrStr | None = None,
425        dialect: DialectType = None,
426        copy: bool = True,
427        **opts: Unpack[ParserNoDialectArgs],
428    ) -> Update:
429        """
430        Set the FROM expression.
431
432        Example:
433            >>> Update().table("my_table").set_("x = 1").from_("baz").sql()
434            'UPDATE my_table SET x = 1 FROM baz'
435
436        Args:
437            expression : the SQL code strings to parse.
438                If a `From` instance is passed, this is used as-is.
439                If another `Expr` instance is passed, it will be wrapped in a `From`.
440                If nothing is passed in then a from is not applied to the expression
441            dialect: the dialect used to parse the input expression.
442            copy: if `False`, modify this expression instance in-place.
443            opts: other options to use to parse the input expressions.
444
445        Returns:
446            The modified Update expression.
447        """
448        if not expression:
449            return maybe_copy(self, copy)
450
451        return _apply_builder(
452            expression=expression,
453            instance=self,
454            arg="from_",
455            into=From,
456            prefix="FROM",
457            dialect=dialect,
458            copy=copy,
459            **opts,
460        )
461
462    def with_(
463        self,
464        alias: ExpOrStr,
465        as_: ExpOrStr,
466        recursive: bool | None = None,
467        materialized: bool | None = None,
468        append: bool = True,
469        dialect: DialectType = None,
470        copy: bool = True,
471        **opts: Unpack[ParserNoDialectArgs],
472    ) -> Update:
473        """
474        Append to or set the common table expressions.
475
476        Example:
477            >>> Update().table("my_table").set_("x = 1").from_("baz").with_("baz", "SELECT id FROM foo").sql()
478            'WITH baz AS (SELECT id FROM foo) UPDATE my_table SET x = 1 FROM baz'
479
480        Args:
481            alias: the SQL code string to parse as the table name.
482                If an `Expr` instance is passed, this is used as-is.
483            as_: the SQL code string to parse as the table expression.
484                If an `Expr` instance is passed, it will be used as-is.
485            recursive: set the RECURSIVE part of the expression. Defaults to `False`.
486            materialized: set the MATERIALIZED part of the expression.
487            append: if `True`, add to any existing expressions.
488                Otherwise, this resets the expressions.
489            dialect: the dialect used to parse the input expression.
490            copy: if `False`, modify this expression instance in-place.
491            opts: other options to use to parse the input expressions.
492
493        Returns:
494            The modified expression.
495        """
496        return _apply_cte_builder(
497            self,
498            alias,
499            as_,
500            recursive=recursive,
501            materialized=materialized,
502            append=append,
503            dialect=dialect,
504            copy=copy,
505            **opts,
506        )
507
508
509class Merge(Expression, DML):
510    arg_types = {
511        "this": True,
512        "using": True,
513        "on": False,
514        "using_cond": False,
515        "whens": True,
516        "with_": False,
517        "returning": False,
518    }
519
520
521class When(Expression):
522    arg_types = {"matched": True, "source": False, "condition": False, "then": True}
523
524
525class Whens(Expression):
526    """Wraps around one or more WHEN [NOT] MATCHED [...] clauses."""
527
528    arg_types = {"expressions": True}