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    }
216
217    def with_(
218        self,
219        alias: ExpOrStr,
220        as_: ExpOrStr,
221        recursive: bool | None = None,
222        materialized: bool | None = None,
223        append: bool = True,
224        dialect: DialectType = None,
225        copy: bool = True,
226        **opts: Unpack[ParserNoDialectArgs],
227    ) -> Insert:
228        """
229        Append to or set the common table expressions.
230
231        Example:
232            >>> import sqlglot
233            >>> sqlglot.parse_one("INSERT INTO t SELECT x FROM cte").with_("cte", as_="SELECT * FROM tbl").sql()
234            'WITH cte AS (SELECT * FROM tbl) INSERT INTO t SELECT x FROM cte'
235
236        Args:
237            alias: the SQL code string to parse as the table name.
238                If an `Expr` instance is passed, this is used as-is.
239            as_: the SQL code string to parse as the table expression.
240                If an `Expr` instance is passed, it will be used as-is.
241            recursive: set the RECURSIVE part of the expression. Defaults to `False`.
242            materialized: set the MATERIALIZED part of the expression.
243            append: if `True`, add to any existing expressions.
244                Otherwise, this resets the expressions.
245            dialect: the dialect used to parse the input expression.
246            copy: if `False`, modify this expression instance in-place.
247            opts: other options to use to parse the input expressions.
248
249        Returns:
250            The modified expression.
251        """
252        return _apply_cte_builder(
253            self,
254            alias,
255            as_,
256            recursive=recursive,
257            materialized=materialized,
258            append=append,
259            dialect=dialect,
260            copy=copy,
261            **opts,
262        )
263
264
265class OnConflict(Expression):
266    arg_types = {
267        "duplicate": False,
268        "expressions": False,
269        "action": False,
270        "conflict_keys": False,
271        "index_predicate": False,
272        "constraint": False,
273        "where": False,
274    }
275
276
277class Returning(Expression):
278    arg_types = {"expressions": True, "into": False}
279
280
281class LoadData(Expression):
282    arg_types = {
283        "this": True,
284        "local": False,
285        "overwrite": False,
286        "inpath": False,
287        "files": False,
288        "partition": False,
289        "input_format": False,
290        "serde": False,
291    }
292
293
294class Update(Expression, DML):
295    arg_types = {
296        "with_": False,
297        "this": False,
298        "expressions": False,
299        "from_": False,
300        "where": False,
301        "returning": False,
302        "order": False,
303        "limit": False,
304        "options": False,
305        "hint": False,
306    }
307
308    def table(
309        self,
310        expression: ExpOrStr,
311        dialect: DialectType = None,
312        copy: bool = True,
313        **opts: Unpack[ParserNoDialectArgs],
314    ) -> Update:
315        """
316        Set the table to update.
317
318        Example:
319            >>> Update().table("my_table").set_("x = 1").sql()
320            'UPDATE my_table SET x = 1'
321
322        Args:
323            expression : the SQL code strings to parse.
324                If a `Table` instance is passed, this is used as-is.
325                If another `Expr` instance is passed, it will be wrapped in a `Table`.
326            dialect: the dialect used to parse the input expression.
327            copy: if `False`, modify this expression instance in-place.
328            opts: other options to use to parse the input expressions.
329
330        Returns:
331            The modified Update expression.
332        """
333        return _apply_builder(
334            expression=expression,
335            instance=self,
336            arg="this",
337            into=Table,
338            prefix=None,
339            dialect=dialect,
340            copy=copy,
341            **opts,
342        )
343
344    def set_(
345        self,
346        *expressions: ExpOrStr,
347        append: bool = True,
348        dialect: DialectType = None,
349        copy: bool = True,
350        **opts: Unpack[ParserNoDialectArgs],
351    ) -> Update:
352        """
353        Append to or set the SET expressions.
354
355        Example:
356            >>> Update().table("my_table").set_("x = 1").sql()
357            'UPDATE my_table SET x = 1'
358
359        Args:
360            *expressions: the SQL code strings to parse.
361                If `Expr` instance(s) are passed, they will be used as-is.
362                Multiple expressions are combined with a comma.
363            append: if `True`, add the new expressions to any existing SET expressions.
364                Otherwise, this resets the expressions.
365            dialect: the dialect used to parse the input expressions.
366            copy: if `False`, modify this expression instance in-place.
367            opts: other options to use to parse the input expressions.
368        """
369        return _apply_list_builder(
370            *expressions,
371            instance=self,
372            arg="expressions",
373            append=append,
374            into=Expr,
375            prefix=None,
376            dialect=dialect,
377            copy=copy,
378            **opts,
379        )
380
381    def where(
382        self,
383        *expressions: ExpOrStr | None,
384        append: bool = True,
385        dialect: DialectType = None,
386        copy: bool = True,
387        **opts: Unpack[ParserNoDialectArgs],
388    ) -> Update:
389        """
390        Append to or set the WHERE expressions.
391
392        Example:
393            >>> Update().table("tbl").set_("x = 1").where("x = 'a' OR x < 'b'").sql()
394            "UPDATE tbl SET x = 1 WHERE x = 'a' OR x < 'b'"
395
396        Args:
397            *expressions: the SQL code strings to parse.
398                If an `Expr` instance is passed, it will be used as-is.
399                Multiple expressions are combined with an AND operator.
400            append: if `True`, AND the new expressions to any existing expression.
401                Otherwise, this resets the expression.
402            dialect: the dialect used to parse the input expressions.
403            copy: if `False`, modify this expression instance in-place.
404            opts: other options to use to parse the input expressions.
405
406        Returns:
407            Update: the modified expression.
408        """
409        return _apply_conjunction_builder(
410            *expressions,
411            instance=self,
412            arg="where",
413            append=append,
414            into=Where,
415            dialect=dialect,
416            copy=copy,
417            **opts,
418        )
419
420    def from_(
421        self,
422        expression: ExpOrStr | None = None,
423        dialect: DialectType = None,
424        copy: bool = True,
425        **opts: Unpack[ParserNoDialectArgs],
426    ) -> Update:
427        """
428        Set the FROM expression.
429
430        Example:
431            >>> Update().table("my_table").set_("x = 1").from_("baz").sql()
432            'UPDATE my_table SET x = 1 FROM baz'
433
434        Args:
435            expression : the SQL code strings to parse.
436                If a `From` instance is passed, this is used as-is.
437                If another `Expr` instance is passed, it will be wrapped in a `From`.
438                If nothing is passed in then a from is not applied to the expression
439            dialect: the dialect used to parse the input expression.
440            copy: if `False`, modify this expression instance in-place.
441            opts: other options to use to parse the input expressions.
442
443        Returns:
444            The modified Update expression.
445        """
446        if not expression:
447            return maybe_copy(self, copy)
448
449        return _apply_builder(
450            expression=expression,
451            instance=self,
452            arg="from_",
453            into=From,
454            prefix="FROM",
455            dialect=dialect,
456            copy=copy,
457            **opts,
458        )
459
460    def with_(
461        self,
462        alias: ExpOrStr,
463        as_: ExpOrStr,
464        recursive: bool | None = None,
465        materialized: bool | None = None,
466        append: bool = True,
467        dialect: DialectType = None,
468        copy: bool = True,
469        **opts: Unpack[ParserNoDialectArgs],
470    ) -> Update:
471        """
472        Append to or set the common table expressions.
473
474        Example:
475            >>> Update().table("my_table").set_("x = 1").from_("baz").with_("baz", "SELECT id FROM foo").sql()
476            'WITH baz AS (SELECT id FROM foo) UPDATE my_table SET x = 1 FROM baz'
477
478        Args:
479            alias: the SQL code string to parse as the table name.
480                If an `Expr` instance is passed, this is used as-is.
481            as_: the SQL code string to parse as the table expression.
482                If an `Expr` instance is passed, it will be used as-is.
483            recursive: set the RECURSIVE part of the expression. Defaults to `False`.
484            materialized: set the MATERIALIZED part of the expression.
485            append: if `True`, add to any existing expressions.
486                Otherwise, this resets the expressions.
487            dialect: the dialect used to parse the input expression.
488            copy: if `False`, modify this expression instance in-place.
489            opts: other options to use to parse the input expressions.
490
491        Returns:
492            The modified expression.
493        """
494        return _apply_cte_builder(
495            self,
496            alias,
497            as_,
498            recursive=recursive,
499            materialized=materialized,
500            append=append,
501            dialect=dialect,
502            copy=copy,
503            **opts,
504        )
505
506
507class Merge(Expression, DML):
508    arg_types = {
509        "this": True,
510        "using": True,
511        "on": False,
512        "using_cond": False,
513        "whens": True,
514        "with_": False,
515        "returning": False,
516    }
517
518
519class When(Expression):
520    arg_types = {"matched": True, "source": False, "condition": False, "then": True}
521
522
523class Whens(Expression):
524    """Wraps around one or more WHEN [NOT] MATCHED [...] clauses."""
525
526    arg_types = {"expressions": True}