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