altimate-code 0.5.1 → 0.5.3

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/__init__.py

@@ -0,0 +1,178 @@
+# ruff: noqa: F401
+"""
+.. include:: ../README.md
+
+----
+"""
+
+from __future__ import annotations
+
+import logging
+import typing as t
+
+from sqlglot import expressions as exp
+from sqlglot.dialects.dialect import Dialect as Dialect, Dialects as Dialects
+from sqlglot.diff import diff as diff
+from sqlglot.errors import (
+    ErrorLevel as ErrorLevel,
+    ParseError as ParseError,
+    TokenError as TokenError,
+    UnsupportedError as UnsupportedError,
+)
+from sqlglot.expressions import (
+    Expression as Expression,
+    alias_ as alias,
+    and_ as and_,
+    case as case,
+    cast as cast,
+    column as column,
+    condition as condition,
+    delete as delete,
+    except_ as except_,
+    from_ as from_,
+    func as func,
+    insert as insert,
+    intersect as intersect,
+    maybe_parse as maybe_parse,
+    merge as merge,
+    not_ as not_,
+    or_ as or_,
+    select as select,
+    subquery as subquery,
+    table_ as table,
+    to_column as to_column,
+    to_identifier as to_identifier,
+    to_table as to_table,
+    union as union,
+)
+from sqlglot.generator import Generator as Generator
+from sqlglot.parser import Parser as Parser
+from sqlglot.schema import MappingSchema as MappingSchema, Schema as Schema
+from sqlglot.tokens import Token as Token, Tokenizer as Tokenizer, TokenType as TokenType
+
+if t.TYPE_CHECKING:
+    from sqlglot._typing import E
+    from sqlglot.dialects.dialect import DialectType as DialectType
+
+logger = logging.getLogger("sqlglot")
+
+
+try:
+    from sqlglot._version import __version__, __version_tuple__
+except ImportError:
+    logger.error(
+        "Unable to set __version__, run `pip install -e .` or `python setup.py develop` first."
+    )
+
+
+pretty = False
+"""Whether to format generated SQL by default."""
+
+
+def tokenize(sql: str, read: DialectType = None, dialect: DialectType = None) -> t.List[Token]:
+    """
+    Tokenizes the given SQL string.
+
+    Args:
+        sql: the SQL code string to tokenize.
+        read: the SQL dialect to apply during tokenizing (eg. "spark", "hive", "presto", "mysql").
+        dialect: the SQL dialect (alias for read).
+
+    Returns:
+        The resulting list of tokens.
+    """
+    return Dialect.get_or_raise(read or dialect).tokenize(sql)
+
+
+def parse(
+    sql: str, read: DialectType = None, dialect: DialectType = None, **opts
+) -> t.List[t.Optional[Expression]]:
+    """
+    Parses the given SQL string into a collection of syntax trees, one per parsed SQL statement.
+
+    Args:
+        sql: the SQL code string to parse.
+        read: the SQL dialect to apply during parsing (eg. "spark", "hive", "presto", "mysql").
+        dialect: the SQL dialect (alias for read).
+        **opts: other `sqlglot.parser.Parser` options.
+
+    Returns:
+        The resulting syntax tree collection.
+    """
+    return Dialect.get_or_raise(read or dialect).parse(sql, **opts)
+
+
+@t.overload
+def parse_one(sql: str, *, into: t.Type[E], **opts) -> E: ...
+
+
+@t.overload
+def parse_one(sql: str, **opts) -> Expression: ...
+
+
+def parse_one(
+    sql: str,
+    read: DialectType = None,
+    dialect: DialectType = None,
+    into: t.Optional[exp.IntoType] = None,
+    **opts,
+) -> Expression:
+    """
+    Parses the given SQL string and returns a syntax tree for the first parsed SQL statement.
+
+    Args:
+        sql: the SQL code string to parse.
+        read: the SQL dialect to apply during parsing (eg. "spark", "hive", "presto", "mysql").
+        dialect: the SQL dialect (alias for read)
+        into: the SQLGlot Expression to parse into.
+        **opts: other `sqlglot.parser.Parser` options.
+
+    Returns:
+        The syntax tree for the first parsed statement.
+    """
+
+    dialect = Dialect.get_or_raise(read or dialect)
+
+    if into:
+        result = dialect.parse_into(into, sql, **opts)
+    else:
+        result = dialect.parse(sql, **opts)
+
+    for expression in result:
+        if not expression:
+            raise ParseError(f"No expression was parsed from '{sql}'")
+        return expression
+    else:
+        raise ParseError(f"No expression was parsed from '{sql}'")
+
+
+def transpile(
+    sql: str,
+    read: DialectType = None,
+    write: DialectType = None,
+    identity: bool = True,
+    error_level: t.Optional[ErrorLevel] = None,
+    **opts,
+) -> t.List[str]:
+    """
+    Parses the given SQL string in accordance with the source dialect and returns a list of SQL strings transformed
+    to conform to the target dialect. Each string in the returned list represents a single transformed SQL statement.
+
+    Args:
+        sql: the SQL code string to transpile.
+        read: the source dialect used to parse the input string (eg. "spark", "hive", "presto", "mysql").
+        write: the target dialect into which the input should be transformed (eg. "spark", "hive", "presto", "mysql").
+        identity: if set to `True` and if the target dialect is not specified the source dialect will be used as both:
+            the source and the target dialect.
+        error_level: the desired error level of the parser.
+        **opts: other `sqlglot.generator.Generator` options.
+
+    Returns:
+        The list of transpiled SQL statements.
+    """
+    write = (read if write is None else write) if identity else write
+    write = Dialect.get_or_raise(write)
+    return [
+        write.generate(expression, copy=False, **opts) if expression else ""
+        for expression in parse(sql, read, error_level=error_level)
+    ]

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/__main__.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import argparse
+import sys
+import typing as t
+
+import sqlglot
+
+parser = argparse.ArgumentParser(description="Transpile SQL")
+parser.add_argument(
+    "sql",
+    metavar="sql",
+    type=str,
+    help="SQL statement(s) to transpile, or - to parse stdin.",
+)
+parser.add_argument(
+    "--read",
+    dest="read",
+    type=str,
+    default=None,
+    help="Dialect to read default is generic",
+)
+parser.add_argument(
+    "--write",
+    dest="write",
+    type=str,
+    default=None,
+    help="Dialect to write default is generic",
+)
+parser.add_argument(
+    "--no-identify",
+    dest="identify",
+    action="store_false",
+    help="Don't auto identify fields",
+)
+parser.add_argument(
+    "--no-pretty",
+    dest="pretty",
+    action="store_false",
+    help="Compress sql",
+)
+parser.add_argument(
+    "--parse",
+    dest="parse",
+    action="store_true",
+    help="Parse and return the expression tree",
+)
+parser.add_argument(
+    "--tokenize",
+    dest="tokenize",
+    action="store_true",
+    help="Tokenize and return the tokens list",
+)
+parser.add_argument(
+    "--error-level",
+    dest="error_level",
+    type=str,
+    default="IMMEDIATE",
+    help="IGNORE, WARN, RAISE, IMMEDIATE (default)",
+)
+parser.add_argument(
+    "--version",
+    action="version",
+    version=sqlglot.__version__,
+    help="Display the SQLGlot version",
+)
+
+
+args = parser.parse_args()
+error_level = sqlglot.ErrorLevel[args.error_level.upper()]
+
+sql = sys.stdin.read() if args.sql == "-" else args.sql
+
+if args.parse:
+    objs: t.Union[t.List[str], t.List[sqlglot.tokens.Token]] = [
+        repr(expression)
+        for expression in sqlglot.parse(
+            sql,
+            read=args.read,
+            error_level=error_level,
+        )
+    ]
+elif args.tokenize:
+    objs = sqlglot.Dialect.get_or_raise(args.read).tokenize(sql)
+else:
+    objs = sqlglot.transpile(
+        sql,
+        read=args.read,
+        write=args.write,
+        identify=args.identify,
+        pretty=args.pretty,
+        error_level=error_level,
+    )
+
+for obj in objs:
+    print(obj)

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/_typing.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import typing as t
+
+import sqlglot
+
+if t.TYPE_CHECKING:
+    from typing_extensions import Literal as Lit  # noqa
+
+# A little hack for backwards compatibility with Python 3.7.
+# For example, we might want a TypeVar for objects that support comparison e.g. SupportsRichComparisonT from typeshed.
+# But Python 3.7 doesn't support Protocols, so we'd also need typing_extensions, which we don't want as a dependency.
+A = t.TypeVar("A", bound=t.Any)
+B = t.TypeVar("B", bound="sqlglot.exp.Binary")
+E = t.TypeVar("E", bound="sqlglot.exp.Expression")
+F = t.TypeVar("F", bound="sqlglot.exp.Func")
+T = t.TypeVar("T")

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dataframe/__init__.py

@@ -0,0 +1,3 @@
+"""
+.. include:: ./README.md
+"""

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dataframe/sql/__init__.py

@@ -0,0 +1,18 @@
+from sqlglot.dataframe.sql.column import Column
+from sqlglot.dataframe.sql.dataframe import DataFrame, DataFrameNaFunctions
+from sqlglot.dataframe.sql.group import GroupedData
+from sqlglot.dataframe.sql.readwriter import DataFrameReader, DataFrameWriter
+from sqlglot.dataframe.sql.session import SparkSession
+from sqlglot.dataframe.sql.window import Window, WindowSpec
+
+__all__ = [
+    "SparkSession",
+    "DataFrame",
+    "GroupedData",
+    "Column",
+    "DataFrameNaFunctions",
+    "Window",
+    "WindowSpec",
+    "DataFrameReader",
+    "DataFrameWriter",
+]

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dataframe/sql/_typing.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import datetime
+import typing as t
+
+from sqlglot import expressions as exp
+
+if t.TYPE_CHECKING:
+    from sqlglot.dataframe.sql.column import Column
+    from sqlglot.dataframe.sql.types import StructType
+
+ColumnLiterals = t.Union[str, float, int, bool, t.List, t.Tuple, datetime.date, datetime.datetime]
+ColumnOrName = t.Union[Column, str]
+ColumnOrLiteral = t.Union[
+    Column, str, float, int, bool, t.List, t.Tuple, datetime.date, datetime.datetime
+]
+SchemaInput = t.Union[str, t.List[str], StructType, t.Dict[str, t.Optional[str]]]
+OutputExpressionContainer = t.Union[exp.Select, exp.Create, exp.Insert]

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dataframe/sql/column.py

@@ -0,0 +1,332 @@
+from __future__ import annotations
+
+import typing as t
+
+import sqlglot as sqlglot
+from sqlglot import expressions as exp
+from sqlglot.dataframe.sql.types import DataType
+from sqlglot.helper import flatten, is_iterable
+
+if t.TYPE_CHECKING:
+    from sqlglot.dataframe.sql._typing import ColumnOrLiteral
+    from sqlglot.dataframe.sql.window import WindowSpec
+
+
+class Column:
+    def __init__(self, expression: t.Optional[t.Union[ColumnOrLiteral, exp.Expression]]):
+        from sqlglot.dataframe.sql.session import SparkSession
+
+        if isinstance(expression, Column):
+            expression = expression.expression  # type: ignore
+        elif expression is None or not isinstance(expression, (str, exp.Expression)):
+            expression = self._lit(expression).expression  # type: ignore
+        elif not isinstance(expression, exp.Column):
+            expression = sqlglot.maybe_parse(expression, dialect=SparkSession().dialect).transform(
+                SparkSession().dialect.normalize_identifier, copy=False
+            )
+        if expression is None:
+            raise ValueError(f"Could not parse {expression}")
+
+        self.expression: exp.Expression = expression  # type: ignore
+
+    def __repr__(self):
+        return repr(self.expression)
+
+    def __hash__(self):
+        return hash(self.expression)
+
+    def __eq__(self, other: ColumnOrLiteral) -> Column:  # type: ignore
+        return self.binary_op(exp.EQ, other)
+
+    def __ne__(self, other: ColumnOrLiteral) -> Column:  # type: ignore
+        return self.binary_op(exp.NEQ, other)
+
+    def __gt__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.GT, other)
+
+    def __ge__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.GTE, other)
+
+    def __lt__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.LT, other)
+
+    def __le__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.LTE, other)
+
+    def __and__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.And, other)
+
+    def __or__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.Or, other)
+
+    def __mod__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.Mod, other)
+
+    def __add__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.Add, other)
+
+    def __sub__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.Sub, other)
+
+    def __mul__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.Mul, other)
+
+    def __truediv__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.Div, other)
+
+    def __div__(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.Div, other)
+
+    def __neg__(self) -> Column:
+        return self.unary_op(exp.Neg)
+
+    def __radd__(self, other: ColumnOrLiteral) -> Column:
+        return self.inverse_binary_op(exp.Add, other)
+
+    def __rsub__(self, other: ColumnOrLiteral) -> Column:
+        return self.inverse_binary_op(exp.Sub, other)
+
+    def __rmul__(self, other: ColumnOrLiteral) -> Column:
+        return self.inverse_binary_op(exp.Mul, other)
+
+    def __rdiv__(self, other: ColumnOrLiteral) -> Column:
+        return self.inverse_binary_op(exp.Div, other)
+
+    def __rtruediv__(self, other: ColumnOrLiteral) -> Column:
+        return self.inverse_binary_op(exp.Div, other)
+
+    def __rmod__(self, other: ColumnOrLiteral) -> Column:
+        return self.inverse_binary_op(exp.Mod, other)
+
+    def __pow__(self, power: ColumnOrLiteral, modulo=None):
+        return Column(exp.Pow(this=self.expression, expression=Column(power).expression))
+
+    def __rpow__(self, power: ColumnOrLiteral):
+        return Column(exp.Pow(this=Column(power).expression, expression=self.expression))
+
+    def __invert__(self):
+        return self.unary_op(exp.Not)
+
+    def __rand__(self, other: ColumnOrLiteral) -> Column:
+        return self.inverse_binary_op(exp.And, other)
+
+    def __ror__(self, other: ColumnOrLiteral) -> Column:
+        return self.inverse_binary_op(exp.Or, other)
+
+    @classmethod
+    def ensure_col(cls, value: t.Optional[t.Union[ColumnOrLiteral, exp.Expression]]) -> Column:
+        return cls(value)
+
+    @classmethod
+    def ensure_cols(cls, args: t.List[t.Union[ColumnOrLiteral, exp.Expression]]) -> t.List[Column]:
+        return [cls.ensure_col(x) if not isinstance(x, Column) else x for x in args]
+
+    @classmethod
+    def _lit(cls, value: ColumnOrLiteral) -> Column:
+        if isinstance(value, dict):
+            columns = [cls._lit(v).alias(k).expression for k, v in value.items()]
+            return cls(exp.Struct(expressions=columns))
+        return cls(exp.convert(value))
+
+    @classmethod
+    def invoke_anonymous_function(
+        cls, column: t.Optional[ColumnOrLiteral], func_name: str, *args: t.Optional[ColumnOrLiteral]
+    ) -> Column:
+        columns = [] if column is None else [cls.ensure_col(column)]
+        column_args = [cls.ensure_col(arg) for arg in args]
+        expressions = [x.expression for x in columns + column_args]
+        new_expression = exp.Anonymous(this=func_name.upper(), expressions=expressions)
+        return Column(new_expression)
+
+    @classmethod
+    def invoke_expression_over_column(
+        cls, column: t.Optional[ColumnOrLiteral], callable_expression: t.Callable, **kwargs
+    ) -> Column:
+        ensured_column = None if column is None else cls.ensure_col(column)
+        ensure_expression_values = {
+            k: [Column.ensure_col(x).expression for x in v]
+            if is_iterable(v)
+            else Column.ensure_col(v).expression
+            for k, v in kwargs.items()
+            if v is not None
+        }
+        new_expression = (
+            callable_expression(**ensure_expression_values)
+            if ensured_column is None
+            else callable_expression(
+                this=ensured_column.column_expression, **ensure_expression_values
+            )
+        )
+        return Column(new_expression)
+
+    def binary_op(self, klass: t.Callable, other: ColumnOrLiteral, **kwargs) -> Column:
+        return Column(
+            klass(this=self.column_expression, expression=Column(other).column_expression, **kwargs)
+        )
+
+    def inverse_binary_op(self, klass: t.Callable, other: ColumnOrLiteral, **kwargs) -> Column:
+        return Column(
+            klass(this=Column(other).column_expression, expression=self.column_expression, **kwargs)
+        )
+
+    def unary_op(self, klass: t.Callable, **kwargs) -> Column:
+        return Column(klass(this=self.column_expression, **kwargs))
+
+    @property
+    def is_alias(self):
+        return isinstance(self.expression, exp.Alias)
+
+    @property
+    def is_column(self):
+        return isinstance(self.expression, exp.Column)
+
+    @property
+    def column_expression(self) -> t.Union[exp.Column, exp.Literal]:
+        return self.expression.unalias()
+
+    @property
+    def alias_or_name(self) -> str:
+        return self.expression.alias_or_name
+
+    @classmethod
+    def ensure_literal(cls, value) -> Column:
+        from sqlglot.dataframe.sql.functions import lit
+
+        if isinstance(value, cls):
+            value = value.expression
+        if not isinstance(value, exp.Literal):
+            return lit(value)
+        return Column(value)
+
+    def copy(self) -> Column:
+        return Column(self.expression.copy())
+
+    def set_table_name(self, table_name: str, copy=False) -> Column:
+        expression = self.expression.copy() if copy else self.expression
+        expression.set("table", exp.to_identifier(table_name))
+        return Column(expression)
+
+    def sql(self, **kwargs) -> str:
+        from sqlglot.dataframe.sql.session import SparkSession
+
+        return self.expression.sql(**{"dialect": SparkSession().dialect, **kwargs})
+
+    def alias(self, name: str) -> Column:
+        new_expression = exp.alias_(self.column_expression, name)
+        return Column(new_expression)
+
+    def asc(self) -> Column:
+        new_expression = exp.Ordered(this=self.column_expression, desc=False, nulls_first=True)
+        return Column(new_expression)
+
+    def desc(self) -> Column:
+        new_expression = exp.Ordered(this=self.column_expression, desc=True, nulls_first=False)
+        return Column(new_expression)
+
+    asc_nulls_first = asc
+
+    def asc_nulls_last(self) -> Column:
+        new_expression = exp.Ordered(this=self.column_expression, desc=False, nulls_first=False)
+        return Column(new_expression)
+
+    def desc_nulls_first(self) -> Column:
+        new_expression = exp.Ordered(this=self.column_expression, desc=True, nulls_first=True)
+        return Column(new_expression)
+
+    desc_nulls_last = desc
+
+    def when(self, condition: Column, value: t.Any) -> Column:
+        from sqlglot.dataframe.sql.functions import when
+
+        column_with_if = when(condition, value)
+        if not isinstance(self.expression, exp.Case):
+            return column_with_if
+        new_column = self.copy()
+        new_column.expression.args["ifs"].extend(column_with_if.expression.args["ifs"])
+        return new_column
+
+    def otherwise(self, value: t.Any) -> Column:
+        from sqlglot.dataframe.sql.functions import lit
+
+        true_value = value if isinstance(value, Column) else lit(value)
+        new_column = self.copy()
+        new_column.expression.set("default", true_value.column_expression)
+        return new_column
+
+    def isNull(self) -> Column:
+        new_expression = exp.Is(this=self.column_expression, expression=exp.Null())
+        return Column(new_expression)
+
+    def isNotNull(self) -> Column:
+        new_expression = exp.Not(this=exp.Is(this=self.column_expression, expression=exp.Null()))
+        return Column(new_expression)
+
+    def cast(self, dataType: t.Union[str, DataType]) -> Column:
+        """
+        Functionality Difference: PySpark cast accepts a datatype instance of the datatype class
+        Sqlglot doesn't currently replicate this class so it only accepts a string
+        """
+        from sqlglot.dataframe.sql.session import SparkSession
+
+        if isinstance(dataType, DataType):
+            dataType = dataType.simpleString()
+        return Column(exp.cast(self.column_expression, dataType, dialect=SparkSession().dialect))
+
+    def startswith(self, value: t.Union[str, Column]) -> Column:
+        value = self._lit(value) if not isinstance(value, Column) else value
+        return self.invoke_anonymous_function(self, "STARTSWITH", value)
+
+    def endswith(self, value: t.Union[str, Column]) -> Column:
+        value = self._lit(value) if not isinstance(value, Column) else value
+        return self.invoke_anonymous_function(self, "ENDSWITH", value)
+
+    def rlike(self, regexp: str) -> Column:
+        return self.invoke_expression_over_column(
+            column=self, callable_expression=exp.RegexpLike, expression=self._lit(regexp).expression
+        )
+
+    def like(self, other: str):
+        return self.invoke_expression_over_column(
+            self, exp.Like, expression=self._lit(other).expression
+        )
+
+    def ilike(self, other: str):
+        return self.invoke_expression_over_column(
+            self, exp.ILike, expression=self._lit(other).expression
+        )
+
+    def substr(self, startPos: t.Union[int, Column], length: t.Union[int, Column]) -> Column:
+        startPos = self._lit(startPos) if not isinstance(startPos, Column) else startPos
+        length = self._lit(length) if not isinstance(length, Column) else length
+        return Column.invoke_expression_over_column(
+            self, exp.Substring, start=startPos.expression, length=length.expression
+        )
+
+    def isin(self, *cols: t.Union[ColumnOrLiteral, t.Iterable[ColumnOrLiteral]]):
+        columns = flatten(cols) if isinstance(cols[0], (list, set, tuple)) else cols  # type: ignore
+        expressions = [self._lit(x).expression for x in columns]
+        return Column.invoke_expression_over_column(self, exp.In, expressions=expressions)  # type: ignore
+
+    def between(
+        self,
+        lowerBound: t.Union[ColumnOrLiteral],
+        upperBound: t.Union[ColumnOrLiteral],
+    ) -> Column:
+        lower_bound_exp = (
+            self._lit(lowerBound) if not isinstance(lowerBound, Column) else lowerBound
+        )
+        upper_bound_exp = (
+            self._lit(upperBound) if not isinstance(upperBound, Column) else upperBound
+        )
+        return Column(
+            exp.Between(
+                this=self.column_expression,
+                low=lower_bound_exp.expression,
+                high=upper_bound_exp.expression,
+            )
+        )
+
+    def over(self, window: WindowSpec) -> Column:
+        window_expression = window.expression.copy()
+        window_expression.set("this", self.column_expression)
+        return Column(window_expression)