Merging upstream version 21.0.1.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

4 files changed, 31 insertions, 11 deletions

diff --git a/sqlglot/executor/__init__.py b/sqlglot/executor/__init__.py
index 304981b..c8f9148 100644
--- a/sqlglot/executor/__init__.py
+++ b/sqlglot/executor/__init__.py
@@ -10,7 +10,6 @@ import logging
 import time
 import typing as t
 
-from sqlglot import maybe_parse
 from sqlglot.errors import ExecuteError
 from sqlglot.executor.python import PythonExecutor
 from sqlglot.executor.table import Table, ensure_tables
@@ -23,7 +22,6 @@ logger = logging.getLogger("sqlglot")
 
 if t.TYPE_CHECKING:
     from sqlglot.dialects.dialect import DialectType
-    from sqlglot.executor.table import Tables
     from sqlglot.expressions import Expression
     from sqlglot.schema import Schema
 
diff --git a/sqlglot/executor/context.py b/sqlglot/executor/context.py
index d7952c1..e4c4040 100644
--- a/sqlglot/executor/context.py
+++ b/sqlglot/executor/context.py
@@ -44,9 +44,9 @@ class Context:
 
             for other in self.tables.values():
                 if self._table.columns != other.columns:
-                    raise Exception(f"Columns are different.")
+                    raise Exception("Columns are different.")
                 if len(self._table.rows) != len(other.rows):
-                    raise Exception(f"Rows are different.")
+                    raise Exception("Rows are different.")
 
         return self._table
 
diff --git a/sqlglot/executor/env.py b/sqlglot/executor/env.py
index 6c01edc..218a8e0 100644
--- a/sqlglot/executor/env.py
+++ b/sqlglot/executor/env.py
@@ -6,7 +6,7 @@ from functools import wraps
 
 from sqlglot import exp
 from sqlglot.generator import Generator
-from sqlglot.helper import PYTHON_VERSION
+from sqlglot.helper import PYTHON_VERSION, is_int, seq_get
 
 
 class reverse_key:
@@ -143,6 +143,22 @@ def arrayjoin(this, expression, null=None):
     return expression.join(x for x in (x if x is not None else null for x in this) if x is not None)
 
 
+@null_if_any("this", "expression")
+def jsonextract(this, expression):
+    for path_segment in expression:
+        if isinstance(this, dict):
+            this = this.get(path_segment)
+        elif isinstance(this, list) and is_int(path_segment):
+            this = seq_get(this, int(path_segment))
+        else:
+            raise NotImplementedError(f"Unable to extract value for {this} at {path_segment}.")
+
+        if this is None:
+            break
+
+    return this
+
+
 ENV = {
     "exp": exp,
     # aggs
@@ -175,12 +191,12 @@ ENV = {
     "DOT": null_if_any(lambda e, this: e[this]),
     "EQ": null_if_any(lambda this, e: this == e),
     "EXTRACT": null_if_any(lambda this, e: getattr(e, this)),
-    "GETPATH": null_if_any(lambda this, e: this.get(e)),
     "GT": null_if_any(lambda this, e: this > e),
     "GTE": null_if_any(lambda this, e: this >= e),
     "IF": lambda predicate, true, false: true if predicate else false,
     "INTDIV": null_if_any(lambda e, this: e // this),
     "INTERVAL": interval,
+    "JSONEXTRACT": jsonextract,
     "LEFT": null_if_any(lambda this, e: this[:e]),
     "LIKE": null_if_any(
         lambda this, e: bool(re.match(e.replace("_", ".").replace("%", ".*"), this))
diff --git a/sqlglot/executor/python.py b/sqlglot/executor/python.py
index 7ff9608..c0becbe 100644
--- a/sqlglot/executor/python.py
+++ b/sqlglot/executor/python.py
@@ -9,7 +9,7 @@ from sqlglot.errors import ExecuteError
 from sqlglot.executor.context import Context
 from sqlglot.executor.env import ENV
 from sqlglot.executor.table import RowReader, Table
-from sqlglot.helper import csv_reader, subclasses
+from sqlglot.helper import csv_reader, ensure_list, subclasses
 
 
 class PythonExecutor:
@@ -368,7 +368,7 @@ def _rename(self, e):
 
         if isinstance(e, exp.Func) and e.is_var_len_args:
             *head, tail = values
-            return self.func(e.key, *head, *tail)
+            return self.func(e.key, *head, *ensure_list(tail))
 
         return self.func(e.key, *values)
     except Exception as ex:
@@ -429,18 +429,24 @@ class Python(Dialect):
             exp.Between: _rename,
             exp.Boolean: lambda self, e: "True" if e.this else "False",
             exp.Cast: lambda self, e: f"CAST({self.sql(e.this)}, exp.DataType.Type.{e.args['to']})",
-            exp.Column: lambda self, e: f"scope[{self.sql(e, 'table') or None}][{self.sql(e.this)}]",
+            exp.Column: lambda self,
+            e: f"scope[{self.sql(e, 'table') or None}][{self.sql(e.this)}]",
             exp.Concat: lambda self, e: self.func(
                 "SAFECONCAT" if e.args.get("safe") else "CONCAT", *e.expressions
             ),
             exp.Distinct: lambda self, e: f"set({self.sql(e, 'this')})",
             exp.Div: _div_sql,
-            exp.Extract: lambda self, e: f"EXTRACT('{e.name.lower()}', {self.sql(e, 'expression')})",
-            exp.In: lambda self, e: f"{self.sql(e, 'this')} in {{{self.expressions(e, flat=True)}}}",
+            exp.Extract: lambda self,
+            e: f"EXTRACT('{e.name.lower()}', {self.sql(e, 'expression')})",
+            exp.In: lambda self,
+            e: f"{self.sql(e, 'this')} in {{{self.expressions(e, flat=True)}}}",
             exp.Interval: lambda self, e: f"INTERVAL({self.sql(e.this)}, '{self.sql(e.unit)}')",
             exp.Is: lambda self, e: (
                 self.binary(e, "==") if isinstance(e.this, exp.Literal) else self.binary(e, "is")
             ),
+            exp.JSONPath: lambda self, e: f"[{','.join(self.sql(p) for p in e.expressions[1:])}]",
+            exp.JSONPathKey: lambda self, e: f"'{self.sql(e.this)}'",
+            exp.JSONPathSubscript: lambda self, e: f"'{e.this}'",
             exp.Lambda: _lambda_sql,
             exp.Not: lambda self, e: f"not {self.sql(e.this)}",
             exp.Null: lambda *_: "None",