ydb-sqlglot-plugin 0.1.1__tar.gz → 0.2.0__tar.gz

{ydb_sqlglot_plugin-0.1.1 → ydb_sqlglot_plugin-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ydb-sqlglot-plugin
-Version: 0.1.1
+Version: 0.2.0
 Summary: YDB dialect plugin for sqlglot
 Author: YDB Team
 License: Apache-2.0
@@ -22,11 +22,12 @@ Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov>=4.0; extra == "dev"
 Requires-Dist: ydb<4,>=3.28.0; extra == "dev"
+Requires-Dist: ruff>=0.9.0; extra == "dev"
 Dynamic: license-file
 
 # ydb-sqlglot-plugin
 
-YDB dialect plugin for [sqlglot](https://github.com/tobymao/sqlglot) — transpiles SQL from any dialect into YDB/YQL.
+YDB dialect plugin for [sqlglot](https://github.com/tobymao/sqlglot) — bidirectional transpilation between YDB/YQL and any SQL dialect.
 
 ## Installation
 
@@ -41,19 +42,20 @@ After installing the package, the `ydb` dialect is available in sqlglot automati
 ```python
 import sqlglot
 
-# Transpile from any dialect
+# Any dialect → YDB
 result = sqlglot.transpile("SELECT * FROM users WHERE id = 1", read="mysql", write="ydb")[0]
 # → SELECT * FROM `users` WHERE id = 1
 
-# Or parse first, then generate
-query = "SELECT * FROM orders WHERE user_id = 1"
-parsed = sqlglot.parse_one(query, dialect="postgres")
-yql = parsed.sql(dialect="ydb")
+# YDB → any dialect
+result = sqlglot.transpile("$t = (SELECT id FROM users); SELECT * FROM $t AS t", read="ydb", write="postgres")[0]
+# → WITH t AS (SELECT id FROM users) SELECT * FROM t AS t
 ```
 
 ## What the plugin does
 
-### Table names
+### Any SQL → YDB
+
+#### Table names
 
 Database-qualified names are rewritten to the YDB path format and wrapped in backticks:
 
@@ -65,7 +67,7 @@ SELECT * FROM analytics.events
 SELECT * FROM `analytics/events`
 ```
 
-### CTEs → YDB variables
+#### CTEs → YDB variables
 
 ```sql
 -- input
@@ -78,7 +80,7 @@ $active = (SELECT * FROM `users` WHERE status = 'active');
 SELECT * FROM $active AS active
 ```
 
-### Subquery decorrelation
+#### Subquery decorrelation
 
 Correlated subqueries (which YQL does not support) are rewritten as JOINs:
 
@@ -102,6 +104,52 @@ The same rewriting applies to `EXISTS`, `IN (subquery)`, and `ANY/ALL` subquerie
 
 ---
 
+### YDB → any SQL
+
+The plugin parses YDB/YQL back into sqlglot's AST, enabling round-trips, YDB-to-YDB transformations, and transpilation to other dialects.
+
+#### Supported YQL constructs
+
+| Construct | Example |
+|---|---|
+| `$variable` references | `SELECT * FROM $t AS t` |
+| `Module::Function()` | `DateTime::GetYear(ts)` |
+| `DECLARE $p AS Type` | `DECLARE $p AS Int32` |
+| `FLATTEN [LIST\|DICT] BY col` | `FROM t FLATTEN LIST BY col` |
+| `Optional<T>` / `T?` | `CAST(x AS Optional<Utf8>)` |
+| Container types | `CAST(x AS List<Int32>)`, `Dict<Utf8, Int64>`, `Set<Utf8>`, `Tuple<Int32, Utf8>` |
+| `ASSUME ORDER BY` | `SELECT * FROM t ASSUME ORDER BY id` |
+| Named expressions | `$t = (SELECT 1 AS x)` |
+| `PRAGMA` | `PRAGMA AnsiImplicitCrossJoin` |
+
+Table names without backticks are accepted on input; the generator always produces backtick-quoted output.
+
+#### CTEs reassembly
+
+YDB-style named expressions are automatically reassembled into standard `WITH` CTEs when targeting other dialects:
+
+```python
+ydb_sql = "$t = (SELECT 1 AS x); SELECT * FROM $t AS t"
+parse_one(ydb_sql, dialect="ydb").sql(dialect="postgres")
+# → WITH t AS (SELECT 1 AS x) SELECT * FROM t AS t
+```
+
+---
+
+### Column lineage
+
+Because YDB SQL is fully parsed into sqlglot's AST, column-level lineage works out of the box:
+
+```python
+from sqlglot.lineage import lineage
+
+node = lineage("total", "$orders = (SELECT user_id, amount FROM orders); SELECT SUM(amount) AS total FROM $orders AS o", dialect="ydb")
+for dep in node.walk():
+    print(dep.name, "→", dep.source)
+```
+
+---
+
 ## Function reference
 
 Functions below are recognized by sqlglot as standard SQL expressions and translated to their YQL equivalents. Dialect-specific functions that sqlglot does not parse into typed AST nodes are **passed through unchanged** — see [Limitations](#limitations).
@@ -179,25 +227,43 @@ Functions below are recognized by sqlglot as standard SQL expressions and transl
 
 ## Type mapping
 
-| SQL type | YQL type |
+### Standard SQL → YDB
+
+| SQL type | YDB type |
 |---|---|
-| `TINYINT` | `INT8` |
-| `SMALLINT` | `INT16` |
-| `INT` / `INTEGER` | `INT32` |
-| `BIGINT` | `INT64` |
+| `TINYINT` | `Int8` |
+| `SMALLINT` | `Int16` |
+| `INT` / `INTEGER` | `Int32` |
+| `BIGINT` | `Int64` |
 | `FLOAT` | `Float` |
 | `DOUBLE` / `DOUBLE PRECISION` | `Double` |
 | `DECIMAL(p, s)` | `Decimal(p, s)` |
 | `BOOLEAN` / `BIT` | `Uint8` |
 | `TIMESTAMP` | `Timestamp` |
-| `VARCHAR` / `NVARCHAR` / `CHAR` | `Utf8` |
-| `TEXT` / `TINYTEXT` / `MEDIUMTEXT` / `LONGTEXT` | `Utf8` |
-| `BLOB` / `TINYBLOB` / `MEDIUMBLOB` / `LONGBLOB` / `VARBINARY` | `String` |
+| `VARCHAR` / `NVARCHAR` / `CHAR` / `TEXT` | `Utf8` |
+| `BLOB` / `BINARY` / `VARBINARY` | `String` |
+
+### YDB types → standard SQL
+
+| YDB type | Standard SQL | Postgres | ClickHouse |
+|---|---|---|---|
+| `Utf8` | `TEXT` | `TEXT` | `String` |
+| `String` | `BLOB` | `BYTEA` | `String` |
+| `Int32` | `INT` | `INT` | `Int32` |
+| `Int64` | `BIGINT` | `BIGINT` | `Int64` |
+| `Optional<T>` | `T` (nullable) | `T` | `Nullable(T)` |
+| `List<T>` | `LIST<T>` | `LIST<T>` | `Array(T)` |
+| `Dict<K,V>` | `MAP<K,V>` | `MAP<K,V>` | `Map(K,V)` |
+| `Tuple<T1,T2>` | `STRUCT<...>` | `STRUCT<...>` | `Tuple(T1,T2)` |
 
 ---
 
 ## Limitations
 
+### Dialect-specific functions
+
+Functions that sqlglot does not parse into typed AST nodes are passed through unchanged and must be replaced manually. Common examples from ClickHouse: `now()`, `today()`, `parseDateTimeBestEffort()`, `toDate()`, `toFloat64()`, `toString()`, `countDistinct()`, `groupArray()`.
+
 ### Correlated subqueries in DML
 
 Correlated subqueries inside `UPDATE` or `INSERT` statements cannot be automatically decorrelated — YDB does not support them natively, and rewriting requires knowledge of the table's primary key. Rewrite manually using a `$variable`:
@@ -217,6 +283,10 @@ Correlated subqueries inside `SELECT` are handled automatically via JOIN rewriti
 
 `dateDiff('month', a, b)` has no exact equivalent in YDB because months have variable length. Use `DateTime::ShiftMonths` for date arithmetic instead.
 
+### YDB container types in other dialects
+
+`Uint8`/`Uint16`/`Uint32`/`Uint64` and YDB-specific container types (`Struct<...>`, `Variant<...>`, `Enum<...>`) do not have direct equivalents in standard SQL and are passed through as-is when targeting other dialects.
+
 ---
 
 ## Development

{ydb_sqlglot_plugin-0.1.1 → ydb_sqlglot_plugin-0.2.0}/README.md

@@ -1,6 +1,6 @@
 # ydb-sqlglot-plugin
 
-YDB dialect plugin for [sqlglot](https://github.com/tobymao/sqlglot) — transpiles SQL from any dialect into YDB/YQL.
+YDB dialect plugin for [sqlglot](https://github.com/tobymao/sqlglot) — bidirectional transpilation between YDB/YQL and any SQL dialect.
 
 ## Installation
 
@@ -15,19 +15,20 @@ After installing the package, the `ydb` dialect is available in sqlglot automati
 ```python
 import sqlglot
 
-# Transpile from any dialect
+# Any dialect → YDB
 result = sqlglot.transpile("SELECT * FROM users WHERE id = 1", read="mysql", write="ydb")[0]
 # → SELECT * FROM `users` WHERE id = 1
 
-# Or parse first, then generate
-query = "SELECT * FROM orders WHERE user_id = 1"
-parsed = sqlglot.parse_one(query, dialect="postgres")
-yql = parsed.sql(dialect="ydb")
+# YDB → any dialect
+result = sqlglot.transpile("$t = (SELECT id FROM users); SELECT * FROM $t AS t", read="ydb", write="postgres")[0]
+# → WITH t AS (SELECT id FROM users) SELECT * FROM t AS t
 ```
 
 ## What the plugin does
 
-### Table names
+### Any SQL → YDB
+
+#### Table names
 
 Database-qualified names are rewritten to the YDB path format and wrapped in backticks:
 
@@ -39,7 +40,7 @@ SELECT * FROM analytics.events
 SELECT * FROM `analytics/events`
 ```
 
-### CTEs → YDB variables
+#### CTEs → YDB variables
 
 ```sql
 -- input
@@ -52,7 +53,7 @@ $active = (SELECT * FROM `users` WHERE status = 'active');
 SELECT * FROM $active AS active
 ```
 
-### Subquery decorrelation
+#### Subquery decorrelation
 
 Correlated subqueries (which YQL does not support) are rewritten as JOINs:
 
@@ -76,6 +77,52 @@ The same rewriting applies to `EXISTS`, `IN (subquery)`, and `ANY/ALL` subquerie
 
 ---
 
+### YDB → any SQL
+
+The plugin parses YDB/YQL back into sqlglot's AST, enabling round-trips, YDB-to-YDB transformations, and transpilation to other dialects.
+
+#### Supported YQL constructs
+
+| Construct | Example |
+|---|---|
+| `$variable` references | `SELECT * FROM $t AS t` |
+| `Module::Function()` | `DateTime::GetYear(ts)` |
+| `DECLARE $p AS Type` | `DECLARE $p AS Int32` |
+| `FLATTEN [LIST\|DICT] BY col` | `FROM t FLATTEN LIST BY col` |
+| `Optional<T>` / `T?` | `CAST(x AS Optional<Utf8>)` |
+| Container types | `CAST(x AS List<Int32>)`, `Dict<Utf8, Int64>`, `Set<Utf8>`, `Tuple<Int32, Utf8>` |
+| `ASSUME ORDER BY` | `SELECT * FROM t ASSUME ORDER BY id` |
+| Named expressions | `$t = (SELECT 1 AS x)` |
+| `PRAGMA` | `PRAGMA AnsiImplicitCrossJoin` |
+
+Table names without backticks are accepted on input; the generator always produces backtick-quoted output.
+
+#### CTEs reassembly
+
+YDB-style named expressions are automatically reassembled into standard `WITH` CTEs when targeting other dialects:
+
+```python
+ydb_sql = "$t = (SELECT 1 AS x); SELECT * FROM $t AS t"
+parse_one(ydb_sql, dialect="ydb").sql(dialect="postgres")
+# → WITH t AS (SELECT 1 AS x) SELECT * FROM t AS t
+```
+
+---
+
+### Column lineage
+
+Because YDB SQL is fully parsed into sqlglot's AST, column-level lineage works out of the box:
+
+```python
+from sqlglot.lineage import lineage
+
+node = lineage("total", "$orders = (SELECT user_id, amount FROM orders); SELECT SUM(amount) AS total FROM $orders AS o", dialect="ydb")
+for dep in node.walk():
+    print(dep.name, "→", dep.source)
+```
+
+---
+
 ## Function reference
 
 Functions below are recognized by sqlglot as standard SQL expressions and translated to their YQL equivalents. Dialect-specific functions that sqlglot does not parse into typed AST nodes are **passed through unchanged** — see [Limitations](#limitations).
@@ -153,25 +200,43 @@ Functions below are recognized by sqlglot as standard SQL expressions and transl
 
 ## Type mapping
 
-| SQL type | YQL type |
+### Standard SQL → YDB
+
+| SQL type | YDB type |
 |---|---|
-| `TINYINT` | `INT8` |
-| `SMALLINT` | `INT16` |
-| `INT` / `INTEGER` | `INT32` |
-| `BIGINT` | `INT64` |
+| `TINYINT` | `Int8` |
+| `SMALLINT` | `Int16` |
+| `INT` / `INTEGER` | `Int32` |
+| `BIGINT` | `Int64` |
 | `FLOAT` | `Float` |
 | `DOUBLE` / `DOUBLE PRECISION` | `Double` |
 | `DECIMAL(p, s)` | `Decimal(p, s)` |
 | `BOOLEAN` / `BIT` | `Uint8` |
 | `TIMESTAMP` | `Timestamp` |
-| `VARCHAR` / `NVARCHAR` / `CHAR` | `Utf8` |
-| `TEXT` / `TINYTEXT` / `MEDIUMTEXT` / `LONGTEXT` | `Utf8` |
-| `BLOB` / `TINYBLOB` / `MEDIUMBLOB` / `LONGBLOB` / `VARBINARY` | `String` |
+| `VARCHAR` / `NVARCHAR` / `CHAR` / `TEXT` | `Utf8` |
+| `BLOB` / `BINARY` / `VARBINARY` | `String` |
+
+### YDB types → standard SQL
+
+| YDB type | Standard SQL | Postgres | ClickHouse |
+|---|---|---|---|
+| `Utf8` | `TEXT` | `TEXT` | `String` |
+| `String` | `BLOB` | `BYTEA` | `String` |
+| `Int32` | `INT` | `INT` | `Int32` |
+| `Int64` | `BIGINT` | `BIGINT` | `Int64` |
+| `Optional<T>` | `T` (nullable) | `T` | `Nullable(T)` |
+| `List<T>` | `LIST<T>` | `LIST<T>` | `Array(T)` |
+| `Dict<K,V>` | `MAP<K,V>` | `MAP<K,V>` | `Map(K,V)` |
+| `Tuple<T1,T2>` | `STRUCT<...>` | `STRUCT<...>` | `Tuple(T1,T2)` |
 
 ---
 
 ## Limitations
 
+### Dialect-specific functions
+
+Functions that sqlglot does not parse into typed AST nodes are passed through unchanged and must be replaced manually. Common examples from ClickHouse: `now()`, `today()`, `parseDateTimeBestEffort()`, `toDate()`, `toFloat64()`, `toString()`, `countDistinct()`, `groupArray()`.
+
 ### Correlated subqueries in DML
 
 Correlated subqueries inside `UPDATE` or `INSERT` statements cannot be automatically decorrelated — YDB does not support them natively, and rewriting requires knowledge of the table's primary key. Rewrite manually using a `$variable`:
@@ -191,6 +256,10 @@ Correlated subqueries inside `SELECT` are handled automatically via JOIN rewriti
 
 `dateDiff('month', a, b)` has no exact equivalent in YDB because months have variable length. Use `DateTime::ShiftMonths` for date arithmetic instead.
 
+### YDB container types in other dialects
+
+`Uint8`/`Uint16`/`Uint32`/`Uint64` and YDB-specific container types (`Struct<...>`, `Variant<...>`, `Enum<...>`) do not have direct equivalents in standard SQL and are passed through as-is when targeting other dialects.
+
 ---
 
 ## Development

{ydb_sqlglot_plugin-0.1.1 → ydb_sqlglot_plugin-0.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ydb-sqlglot-plugin"
-version = "0.1.1"  # AUTOVERSION
+version = "0.2.0"  # AUTOVERSION
 description = "YDB dialect plugin for sqlglot"
 readme = "README.md"
 license = {text = "Apache-2.0"}
@@ -31,6 +31,7 @@ dev = [
     "pytest>=7.0",
     "pytest-cov>=4.0",
     "ydb>=3.28.0,<4",
+    "ruff>=0.9.0",
 ]
 
 [project.urls]
@@ -48,3 +49,11 @@ include = ["ydb_sqlglot*"]
 testpaths = ["tests"]
 python_files = ["test_*.py"]
 python_functions = ["test_*"]
+
+[tool.ruff]
+target-version = "py39"
+line-length = 120
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "W"]
+ignore = ["E501"]

{ydb_sqlglot_plugin-0.1.1 → ydb_sqlglot_plugin-0.2.0}/ydb_sqlglot/__init__.py

@@ -1,7 +1,6 @@
 """YDB SQL dialect plugin for sqlglot."""
 
-from .ydb import YDB
 from .version import VERSION
-
+from .ydb import YDB as YDB
 
 __version__ = VERSION

ydb_sqlglot_plugin-0.2.0/ydb_sqlglot/version.py

@@ -0,0 +1 @@
+VERSION = "0.2.0"

{ydb_sqlglot_plugin-0.1.1 → ydb_sqlglot_plugin-0.2.0}/ydb_sqlglot/ydb.py

@@ -1,16 +1,14 @@
 import re
 import typing as t
 
-from sqlglot import exp, tokens, generator, transforms, TokenType, parser, Generator
+from sqlglot import Generator, TokenType, exp, generator, parser, tokens, transforms
+from sqlglot.dialects.dialect import Dialect, NormalizationStrategy, concat_to_dpipe_sql, unit_to_var
 from sqlglot.errors import UnsupportedError
 from sqlglot.expressions import Expression
-from sqlglot.dialects.dialect import Dialect, unit_to_var
-from sqlglot.dialects.dialect import NormalizationStrategy, concat_to_dpipe_sql
-from sqlglot.helper import name_sequence, seq_get, flatten
+from sqlglot.helper import flatten, name_sequence, seq_get
+from sqlglot.optimizer.scope import ScopeType, find_in_scope, traverse_scope
 from sqlglot.optimizer.simplify import simplify
-from sqlglot.transforms import move_ctes_to_top_level
-from sqlglot.optimizer.scope import find_in_scope, ScopeType, traverse_scope
-from sqlglot.transforms import eliminate_join_marks
+from sqlglot.transforms import eliminate_join_marks, move_ctes_to_top_level
 
 JOIN_ATTRS = ("on", "side", "kind", "using", "method")
 
@@ -29,21 +27,6 @@ def table_names_to_lower_case(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
-def make_db_name_lower(expression: exp.Expression) -> exp.Expression:
-    """
-    Converts all database names to uppercase
-    Args:
-        expression: The SQL expression to modify
-    Returns:
-        Modified expression with uppercase database names
-    """
-    for table in expression.find_all(exp.Table):
-        if table.db:
-            table.set("db", table.db.lower())
-
-    return expression
-
-
 def make_db_name_lower(expression: exp.Expression) -> exp.Expression:
     """
     Converts all database names to uppercase
@@ -440,6 +423,110 @@ def _has_implicit_cross_join(expression: exp.Expression) -> bool:
     return False
 
 
+class FlattenBy(exp.Expression):
+    """YDB-specific FLATTEN [LIST|DICT] BY clause on a table reference."""
+    arg_types = {"this": True, "expressions": True, "kind": False}
+
+
+class AssumeOrderBy(exp.Expression):
+    """YDB-specific ASSUME ORDER BY hint (data is pre-sorted, skip sort)."""
+    arg_types = {"this": True}
+
+
+class YdbTuple(exp.Expression):
+    """YDB Tuple<T1, T2, ...> type — positional unnamed fields."""
+    arg_types = {"expressions": True, "nullable": False}
+
+
+# Container types that use Generic<T, ...> syntax in YDB
+_YDB_GENERIC_TYPES = {
+    "List": exp.DataType.Type.LIST,
+    "Dict": exp.DataType.Type.MAP,
+    "Set": exp.DataType.Type.SET,
+}
+
+
+def _reassemble_ctes(
+    statements: t.List[t.Optional[exp.Expression]],
+) -> t.List[t.Optional[exp.Expression]]:
+    """Convert sequences of YDB named-expression statements into standard WITH CTEs.
+
+    YDB generator emits:  $t = (SELECT ...);  SELECT * FROM $t AS t
+    This function rebuilds:  WITH t AS (SELECT ...) SELECT * FROM t
+
+    so that transpiling YDB output to other dialects produces valid SQL.
+    """
+    result: t.List[t.Optional[exp.Expression]] = []
+    # Keep both the original Alias nodes and the converted CTE nodes
+    pending_aliases: t.List[exp.Alias] = []
+    pending_ctes: t.List[exp.CTE] = []
+    pending_names: t.Set[str] = set()
+
+    def _flush_as_aliases() -> None:
+        result.extend(pending_aliases)
+        pending_aliases.clear()
+        pending_ctes.clear()
+        pending_names.clear()
+
+    for stmt in statements:
+        if (
+            isinstance(stmt, exp.Alias)
+            and isinstance(stmt.args.get("alias"), exp.Identifier)
+            and stmt.alias.startswith("$")
+        ):
+            name = stmt.alias[1:]
+            inner = stmt.this
+            # Unwrap Subquery — CTE.this must be Select, not Subquery
+            if isinstance(inner, exp.Subquery):
+                inner = inner.this
+            # Replace any $prev_cte refs inside this CTE body
+            inner = _replace_param_table_refs(inner, pending_names)
+            pending_aliases.append(stmt)
+            pending_ctes.append(
+                exp.CTE(
+                    this=inner,
+                    alias=exp.TableAlias(this=exp.to_identifier(name)),
+                )
+            )
+            pending_names.add(name)
+        elif pending_ctes and isinstance(stmt, exp.Select):
+            stmt = _replace_param_table_refs(stmt, pending_names)
+            stmt.set("with_", exp.With(expressions=list(pending_ctes)))
+            result.append(stmt)
+            pending_aliases.clear()
+            pending_ctes.clear()
+            pending_names.clear()
+        else:
+            # No following SELECT — keep original Alias form
+            _flush_as_aliases()
+            result.append(stmt)
+
+    # Trailing named exprs without a SELECT — keep as-is
+    _flush_as_aliases()
+    return result
+
+
+def _replace_param_table_refs(
+    tree: exp.Expression, names: t.Set[str]
+) -> exp.Expression:
+    """Replace Table(Parameter(Var("t"))) with Table(Identifier("t")) for CTE names."""
+
+    def _transform(node: exp.Expression) -> exp.Expression:
+        if (
+            isinstance(node, exp.Table)
+            and isinstance(node.this, exp.Parameter)
+            and isinstance(node.this.this, exp.Var)
+            and node.this.this.name in names
+        ):
+            return exp.Table(
+                this=exp.to_identifier(node.this.this.name),
+                alias=node.args.get("alias"),
+            )
+        return node
+
+    return tree.transform(_transform)
+
+
 class YDB(Dialect):
     """
     YDB SQL dialect implementation for sqlglot.
@@ -470,9 +557,18 @@ class YDB(Dialect):
         Defines how the SQL text is broken into tokens.
         """
 
+        KEYWORDS = {
+            **tokens.Tokenizer.KEYWORDS,
+            "DECLARE": TokenType.DECLARE,
+            "UTF8": TokenType.TEXT,       # YDB Utf8 = unicode text = SQL TEXT
+            "STRING": TokenType.BLOB,     # YDB String = bytes = SQL BLOB
+        }
+
         SINGLE_TOKENS = {
             **tokens.Tokenizer.SINGLE_TOKENS,
+            "$": TokenType.PARAMETER,
         }
+        VAR_SINGLE_TOKENS = {"$"}
 
         SUPPORTS_VALUES_DEFAULT = False
         QUOTES = ["'", '"']
@@ -480,6 +576,154 @@ class YDB(Dialect):
         IDENTIFIERS = ["`"]
 
     class Parser(parser.Parser):
+        COLUMN_OPERATORS = {
+            **parser.Parser.COLUMN_OPERATORS,
+            # In YDB :: is a module namespace separator (e.g. DateTime::GetYear),
+            # not a Postgres-style cast. Reparse the right side as a function call.
+            TokenType.DCOLON: lambda self, this, field: (
+                self.expression(
+                    exp.Anonymous(
+                        this=f"{this.name}::{field.name}",
+                        expressions=field.expressions,
+                    )
+                )
+                if isinstance(field, exp.Func)
+                else self.expression(exp.ScopeResolution(this=this, expression=field))
+            ),
+        }
+
+        STATEMENT_PARSERS = {
+            **parser.Parser.STATEMENT_PARSERS,
+            TokenType.DECLARE: lambda self: self._parse_ydb_declare(),
+            TokenType.PARAMETER: lambda self: self._parse_ydb_named_expr(),
+        }
+
+        def parse(self, raw_tokens, sql=None):
+            statements = super().parse(raw_tokens, sql)
+            return _reassemble_ctes(statements)
+
+        def _parse_dcolon(self) -> t.Optional[exp.Expression]:
+            return self._parse_function(anonymous=True) or self._parse_var(any_token=True)
+
+        def _parse_ydb_named_expr(self) -> t.Optional[exp.Expression]:
+            # _match_set already consumed '$', so _index points to the var name.
+            # Retreat one extra step to include '$' when falling back to expression parsing.
+            index = self._index - 1
+            name_var = self._parse_var(any_token=True)
+            if not self._match(TokenType.EQ):
+                # Not an assignment — retreat (including '$') and parse as expression.
+                self._retreat(index)
+                return self._parse_expression()
+            value = self._parse_select() or self._parse_expression()
+            return self.expression(
+                exp.Alias(
+                    this=value,
+                    alias=exp.Identifier(this=f"${name_var.name}"),
+                )
+            )
+
+        def _parse_ydb_declare(self) -> exp.Declare:
+            items = self._parse_csv(self._parse_ydb_declareitem)
+            return self.expression(exp.Declare(expressions=items))
+
+        def _parse_ydb_declareitem(self) -> t.Optional[exp.DeclareItem]:
+            if not self._match(TokenType.PARAMETER):
+                return None
+            name = self._parse_var(any_token=True)
+            if not name:
+                return None
+            self._match(TokenType.ALIAS)
+            kind = self._parse_types()
+            return self.expression(exp.DeclareItem(this=name, kind=kind))
+
+        def _parse_types(self, *args, **kwargs) -> t.Optional[exp.Expression]:
+            # YDB generic types use Name<...> syntax; token type varies by keyword status
+            if self._curr and self._next and self._next.token_type == TokenType.LT:
+                name = self._curr.text
+
+                if name == "Optional":
+                    self._advance()  # consume 'Optional'
+                    self._advance()  # consume '<'
+                    inner = self._parse_types(*args, **kwargs)
+                    self._match(TokenType.GT)
+                    if inner:
+                        inner.set("nullable", True)
+                    return inner
+
+                if name in _YDB_GENERIC_TYPES:
+                    self._advance()  # consume type name
+                    self._advance()  # consume '<'
+                    type_args = self._parse_csv(
+                        lambda: self._parse_types(*args, **kwargs)
+                    )
+                    self._match(TokenType.GT)
+                    return exp.DataType(
+                        this=_YDB_GENERIC_TYPES[name],
+                        expressions=[a for a in type_args if a],
+                        nested=True,
+                    )
+
+                if name == "Tuple":
+                    self._advance()  # consume 'Tuple'
+                    self._advance()  # consume '<'
+                    type_args = self._parse_csv(
+                        lambda: self._parse_types(*args, **kwargs)
+                    )
+                    self._match(TokenType.GT)
+                    # Represent as STRUCT so other dialects can serialize it.
+                    # kind="tuple" is a YDB-specific marker for the generator to emit Tuple<...>.
+                    return exp.DataType(
+                        this=exp.DataType.Type.STRUCT,
+                        expressions=[
+                            exp.ColumnDef(this=exp.to_identifier(f"_{i}"), kind=a)
+                            for i, a in enumerate(type_args) if a
+                        ],
+                        nested=True,
+                        kind=exp.Var(this="tuple"),
+                    )
+
+            dtype = super()._parse_types(*args, **kwargs)
+            if dtype and self._match(TokenType.PLACEHOLDER):  # T?
+                dtype.set("nullable", True)
+            return dtype
+
+        def _parse_table_alias(self, alias_tokens=None):
+            # Prevent YDB-specific keywords from being consumed as table aliases
+            if self._curr and self._curr.text.upper() in ("FLATTEN", "ASSUME"):
+                # Also check that what follows is a YDB construct, not a regular alias
+                if self._next and (
+                    self._next.text.upper() in ("BY", "LIST", "DICT")
+                    or self._next.token_type == TokenType.ORDER_BY
+                ):
+                    return None
+            return super()._parse_table_alias(alias_tokens=alias_tokens)
+
+        def _parse_query_modifiers(self, this):
+            if (
+                self._curr
+                and self._curr.text.upper() == "ASSUME"
+                and self._next
+                and self._next.token_type == TokenType.ORDER_BY
+            ):
+                self._advance()  # consume ASSUME
+                _, order = self.QUERY_MODIFIER_PARSERS[TokenType.ORDER_BY](self)
+                if order and this:
+                    this.set("order", self.expression(AssumeOrderBy(this=order)))
+            return super()._parse_query_modifiers(this)
+
+        def _parse_table(self, *args, **kwargs) -> t.Optional[exp.Expression]:
+            table = super()._parse_table(*args, **kwargs)
+            if table and self._curr and self._curr.text.upper() == "FLATTEN":
+                self._advance()
+                kind: t.Optional[str] = None
+                if self._curr and self._curr.text.upper() in ("LIST", "DICT"):
+                    kind = self._curr.text.upper()
+                    self._advance()
+                self._match_text_seq("BY")
+                cols = self._parse_csv(self._parse_column)
+                return self.expression(FlattenBy(this=table, expressions=cols, kind=kind))
+            return table
+
         def _parse_struct_types(self, type_required=True) -> t.Optional[exp.Expression]:
             if not self._curr:
                 return None
@@ -559,6 +803,8 @@ class YDB(Dialect):
         Responsible for translating SQL AST back to SQL text with YDB-specific syntax.
         """
 
+        PARAMETER_TOKEN = "$"
+
         SUPPORTS_VALUES_DEFAULT = False
         NORMALIZATION_STRATEGY = NormalizationStrategy.CASE_SENSITIVE
         JOIN_HINTS = False
@@ -574,7 +820,7 @@ class YDB(Dialect):
         JSON_KEY_VALUE_PAIR_SEP = ","
         VARCHAR_REQUIRES_SIZE = False
         CAN_IMPLEMENT_ARRAY_ANY = True
-        STRUCT_DELIMITER = ("<|", "|>")
+        STRUCT_DELIMITER = ("<", ">")
         NULL_ORDERING_SUPPORTED: t.Optional[bool] = False
         NULL_ORDERING = None
         MATCHED_BY_SOURCE = False
@@ -637,6 +883,10 @@ class YDB(Dialect):
             Returns:
                 Generated SQL string for the table reference
             """
+            if isinstance(expression.this, exp.Parameter):
+                var = self.sql(expression, "this")
+                alias = f" AS {expression.alias}" if expression.alias else ""
+                return f"{var}{alias}"
             prefix = f"{expression.db}/" if expression.db else ""
             sql = f"`{prefix}{expression.name}`"
 
@@ -663,6 +913,38 @@ class YDB(Dialect):
 
             return is_sql
 
+        def scoperesolution_sql(self, expression: exp.ScopeResolution) -> str:
+            this = self.sql(expression, "this")
+            expr = self.sql(expression, "expression")
+            return f"{this}::{expr}"
+
+        def declareitem_sql(self, expression: exp.DeclareItem) -> str:
+            name = self.sql(expression, "this")
+            kind = self.sql(expression, "kind")
+            return f"${name} AS {kind}"
+
+        def flattenby_sql(self, expression: FlattenBy) -> str:
+            table = self.sql(expression, "this")
+            kind = expression.args.get("kind")
+            kind_str = f" {kind}" if kind else ""
+            cols = self.expressions(expression, flat=True)
+            return f"{table} FLATTEN{kind_str} BY {cols}"
+
+        def assumeorderby_sql(self, expression: AssumeOrderBy) -> str:
+            order = self.sql(expression, "this").lstrip()
+            return self.seg(f"ASSUME {order}")
+
+        def ydbtuple_sql(self, expression: YdbTuple) -> str:
+            inner = ", ".join(self.sql(e) for e in expression.expressions)
+            sql = f"Tuple<{inner}>"
+            return f"Optional<{sql}>" if expression.args.get("nullable") else sql
+
+        def alias_sql(self, expression: exp.Alias) -> str:
+            alias = expression.args.get("alias")
+            if alias and alias.name.startswith("$"):
+                return f"{alias.name} = {self.sql(expression, 'this')}"
+            return super().alias_sql(expression)
+
         def anonymous_sql(self, expression: exp.Anonymous) -> str:
             """
             Generate SQL for Anonymous functions, with special handling for YQL lambda variables.
@@ -799,6 +1081,35 @@ class YDB(Dialect):
             Returns:
                 Generated SQL string for the data type
             """
+            nullable = expression.args.get("nullable")
+
+            # YDB generic container types rendered with <> syntax and correct casing
+            if expression.args.get("nested"):
+                type_value = expression.this
+                # Tuple<...>: STRUCT with kind="tuple" marker
+                if (
+                    type_value == exp.DataType.Type.STRUCT
+                    and isinstance(expression.args.get("kind"), exp.Var)
+                    and expression.args["kind"].name == "tuple"
+                ):
+                    inner = ", ".join(
+                        self.sql(col.args["kind"])
+                        for col in expression.expressions
+                        if isinstance(col, exp.ColumnDef)
+                    )
+                    sql = f"Tuple<{inner}>"
+                    return f"Optional<{sql}>" if nullable else sql
+
+                inner = ", ".join(self.sql(e) for e in expression.expressions)
+                name = {
+                    exp.DataType.Type.LIST: "List",
+                    exp.DataType.Type.MAP: "Dict",
+                    exp.DataType.Type.SET: "Set",
+                }.get(type_value)
+                if name:
+                    sql = f"{name}<{inner}>"
+                    return f"Optional<{sql}>" if nullable else sql
+
             if (
                     expression.is_type(exp.DataType.Type.NVARCHAR)
                     or expression.is_type(exp.DataType.Type.VARCHAR)
@@ -845,7 +1156,10 @@ class YDB(Dialect):
                         exp.DataType.build("float") if size <= 32 else exp.DataType.build("double")
                     )
 
-            return super().datatype_sql(expression)
+            sql = super().datatype_sql(expression)
+            if nullable:
+                sql = f"Optional<{sql}>"
+            return sql
 
         def primarykeycolumnconstraint_sql(self, expression: exp.PrimaryKeyColumnConstraint) -> str:
             """
@@ -2358,10 +2672,10 @@ class YDB(Dialect):
         TYPE_MAPPING = {
             **generator.Generator.TYPE_MAPPING,
             **STRING_TYPE_MAPPING,
-            exp.DataType.Type.TINYINT: "INT8",
-            exp.DataType.Type.SMALLINT: "INT16",
-            exp.DataType.Type.INT: "INT32",
-            exp.DataType.Type.BIGINT: "INT64",
+            exp.DataType.Type.TINYINT: "Int8",
+            exp.DataType.Type.SMALLINT: "Int16",
+            exp.DataType.Type.INT: "Int32",
+            exp.DataType.Type.BIGINT: "Int64",
             exp.DataType.Type.DECIMAL: "Decimal",
             exp.DataType.Type.FLOAT: "Float",
             exp.DataType.Type.DOUBLE: "Double",
@@ -2373,6 +2687,9 @@ class YDB(Dialect):
 
         TRANSFORMS = {
             **generator.Generator.TRANSFORMS,
+            FlattenBy: lambda self, e: self.flattenby_sql(e),
+            AssumeOrderBy: lambda self, e: self.assumeorderby_sql(e),
+            YdbTuple: lambda self, e: self.ydbtuple_sql(e),
             exp.Create: create_sql,
             exp.DefaultColumnConstraint: lambda self, e: "",
             exp.DateTrunc: _date_trunc_sql,

{ydb_sqlglot_plugin-0.1.1 → ydb_sqlglot_plugin-0.2.0}/ydb_sqlglot_plugin.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ydb-sqlglot-plugin
-Version: 0.1.1
+Version: 0.2.0
 Summary: YDB dialect plugin for sqlglot
 Author: YDB Team
 License: Apache-2.0
@@ -22,11 +22,12 @@ Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov>=4.0; extra == "dev"
 Requires-Dist: ydb<4,>=3.28.0; extra == "dev"
+Requires-Dist: ruff>=0.9.0; extra == "dev"
 Dynamic: license-file
 
 # ydb-sqlglot-plugin
 
-YDB dialect plugin for [sqlglot](https://github.com/tobymao/sqlglot) — transpiles SQL from any dialect into YDB/YQL.
+YDB dialect plugin for [sqlglot](https://github.com/tobymao/sqlglot) — bidirectional transpilation between YDB/YQL and any SQL dialect.
 
 ## Installation
 
@@ -41,19 +42,20 @@ After installing the package, the `ydb` dialect is available in sqlglot automati
 ```python
 import sqlglot
 
-# Transpile from any dialect
+# Any dialect → YDB
 result = sqlglot.transpile("SELECT * FROM users WHERE id = 1", read="mysql", write="ydb")[0]
 # → SELECT * FROM `users` WHERE id = 1
 
-# Or parse first, then generate
-query = "SELECT * FROM orders WHERE user_id = 1"
-parsed = sqlglot.parse_one(query, dialect="postgres")
-yql = parsed.sql(dialect="ydb")
+# YDB → any dialect
+result = sqlglot.transpile("$t = (SELECT id FROM users); SELECT * FROM $t AS t", read="ydb", write="postgres")[0]
+# → WITH t AS (SELECT id FROM users) SELECT * FROM t AS t
 ```
 
 ## What the plugin does
 
-### Table names
+### Any SQL → YDB
+
+#### Table names
 
 Database-qualified names are rewritten to the YDB path format and wrapped in backticks:
 
@@ -65,7 +67,7 @@ SELECT * FROM analytics.events
 SELECT * FROM `analytics/events`
 ```
 
-### CTEs → YDB variables
+#### CTEs → YDB variables
 
 ```sql
 -- input
@@ -78,7 +80,7 @@ $active = (SELECT * FROM `users` WHERE status = 'active');
 SELECT * FROM $active AS active
 ```
 
-### Subquery decorrelation
+#### Subquery decorrelation
 
 Correlated subqueries (which YQL does not support) are rewritten as JOINs:
 
@@ -102,6 +104,52 @@ The same rewriting applies to `EXISTS`, `IN (subquery)`, and `ANY/ALL` subquerie
 
 ---
 
+### YDB → any SQL
+
+The plugin parses YDB/YQL back into sqlglot's AST, enabling round-trips, YDB-to-YDB transformations, and transpilation to other dialects.
+
+#### Supported YQL constructs
+
+| Construct | Example |
+|---|---|
+| `$variable` references | `SELECT * FROM $t AS t` |
+| `Module::Function()` | `DateTime::GetYear(ts)` |
+| `DECLARE $p AS Type` | `DECLARE $p AS Int32` |
+| `FLATTEN [LIST\|DICT] BY col` | `FROM t FLATTEN LIST BY col` |
+| `Optional<T>` / `T?` | `CAST(x AS Optional<Utf8>)` |
+| Container types | `CAST(x AS List<Int32>)`, `Dict<Utf8, Int64>`, `Set<Utf8>`, `Tuple<Int32, Utf8>` |
+| `ASSUME ORDER BY` | `SELECT * FROM t ASSUME ORDER BY id` |
+| Named expressions | `$t = (SELECT 1 AS x)` |
+| `PRAGMA` | `PRAGMA AnsiImplicitCrossJoin` |
+
+Table names without backticks are accepted on input; the generator always produces backtick-quoted output.
+
+#### CTEs reassembly
+
+YDB-style named expressions are automatically reassembled into standard `WITH` CTEs when targeting other dialects:
+
+```python
+ydb_sql = "$t = (SELECT 1 AS x); SELECT * FROM $t AS t"
+parse_one(ydb_sql, dialect="ydb").sql(dialect="postgres")
+# → WITH t AS (SELECT 1 AS x) SELECT * FROM t AS t
+```
+
+---
+
+### Column lineage
+
+Because YDB SQL is fully parsed into sqlglot's AST, column-level lineage works out of the box:
+
+```python
+from sqlglot.lineage import lineage
+
+node = lineage("total", "$orders = (SELECT user_id, amount FROM orders); SELECT SUM(amount) AS total FROM $orders AS o", dialect="ydb")
+for dep in node.walk():
+    print(dep.name, "→", dep.source)
+```
+
+---
+
 ## Function reference
 
 Functions below are recognized by sqlglot as standard SQL expressions and translated to their YQL equivalents. Dialect-specific functions that sqlglot does not parse into typed AST nodes are **passed through unchanged** — see [Limitations](#limitations).
@@ -179,25 +227,43 @@ Functions below are recognized by sqlglot as standard SQL expressions and transl
 
 ## Type mapping
 
-| SQL type | YQL type |
+### Standard SQL → YDB
+
+| SQL type | YDB type |
 |---|---|
-| `TINYINT` | `INT8` |
-| `SMALLINT` | `INT16` |
-| `INT` / `INTEGER` | `INT32` |
-| `BIGINT` | `INT64` |
+| `TINYINT` | `Int8` |
+| `SMALLINT` | `Int16` |
+| `INT` / `INTEGER` | `Int32` |
+| `BIGINT` | `Int64` |
 | `FLOAT` | `Float` |
 | `DOUBLE` / `DOUBLE PRECISION` | `Double` |
 | `DECIMAL(p, s)` | `Decimal(p, s)` |
 | `BOOLEAN` / `BIT` | `Uint8` |
 | `TIMESTAMP` | `Timestamp` |
-| `VARCHAR` / `NVARCHAR` / `CHAR` | `Utf8` |
-| `TEXT` / `TINYTEXT` / `MEDIUMTEXT` / `LONGTEXT` | `Utf8` |
-| `BLOB` / `TINYBLOB` / `MEDIUMBLOB` / `LONGBLOB` / `VARBINARY` | `String` |
+| `VARCHAR` / `NVARCHAR` / `CHAR` / `TEXT` | `Utf8` |
+| `BLOB` / `BINARY` / `VARBINARY` | `String` |
+
+### YDB types → standard SQL
+
+| YDB type | Standard SQL | Postgres | ClickHouse |
+|---|---|---|---|
+| `Utf8` | `TEXT` | `TEXT` | `String` |
+| `String` | `BLOB` | `BYTEA` | `String` |
+| `Int32` | `INT` | `INT` | `Int32` |
+| `Int64` | `BIGINT` | `BIGINT` | `Int64` |
+| `Optional<T>` | `T` (nullable) | `T` | `Nullable(T)` |
+| `List<T>` | `LIST<T>` | `LIST<T>` | `Array(T)` |
+| `Dict<K,V>` | `MAP<K,V>` | `MAP<K,V>` | `Map(K,V)` |
+| `Tuple<T1,T2>` | `STRUCT<...>` | `STRUCT<...>` | `Tuple(T1,T2)` |
 
 ---
 
 ## Limitations
 
+### Dialect-specific functions
+
+Functions that sqlglot does not parse into typed AST nodes are passed through unchanged and must be replaced manually. Common examples from ClickHouse: `now()`, `today()`, `parseDateTimeBestEffort()`, `toDate()`, `toFloat64()`, `toString()`, `countDistinct()`, `groupArray()`.
+
 ### Correlated subqueries in DML
 
 Correlated subqueries inside `UPDATE` or `INSERT` statements cannot be automatically decorrelated — YDB does not support them natively, and rewriting requires knowledge of the table's primary key. Rewrite manually using a `$variable`:
@@ -217,6 +283,10 @@ Correlated subqueries inside `SELECT` are handled automatically via JOIN rewriti
 
 `dateDiff('month', a, b)` has no exact equivalent in YDB because months have variable length. Use `DateTime::ShiftMonths` for date arithmetic instead.
 
+### YDB container types in other dialects
+
+`Uint8`/`Uint16`/`Uint32`/`Uint64` and YDB-specific container types (`Struct<...>`, `Variant<...>`, `Enum<...>`) do not have direct equivalents in standard SQL and are passed through as-is when targeting other dialects.
+
 ---
 
 ## Development

{ydb_sqlglot_plugin-0.1.1 → ydb_sqlglot_plugin-0.2.0}/ydb_sqlglot_plugin.egg-info/requires.txt

@@ -4,3 +4,4 @@ sqlglot>=28.6.0
 pytest>=7.0
 pytest-cov>=4.0
 ydb<4,>=3.28.0
+ruff>=0.9.0

ydb_sqlglot_plugin-0.1.1/ydb_sqlglot/version.py

@@ -1 +0,0 @@
-VERSION = "0.1.1"