From e124ff538f512716a4d4779d495c6f4e38e3e072 Mon Sep 17 00:00:00 2001 From: Bartok Date: Mon, 29 Jun 2026 20:11:47 -0600 Subject: [PATCH] feat(wren): add cross-dialect type translation to type_mapping (#2410) --- core/wren/src/wren/type_mapping.py | 54 ++++++++ core/wren/src/wren/utils_cli.py | 67 +++++++++- core/wren/tests/unit/test_type_mapping.py | 151 +++++++++++++++++++++- 3 files changed, 270 insertions(+), 2 deletions(-) diff --git a/core/wren/src/wren/type_mapping.py b/core/wren/src/wren/type_mapping.py index 9a939f8f2..b043d694c 100644 --- a/core/wren/src/wren/type_mapping.py +++ b/core/wren/src/wren/type_mapping.py @@ -14,6 +14,8 @@ Use as a library: # {"column": "id", "raw_type": "int8", "type": "BIGINT"}, # {"column": "name", "raw_type": "character varying", "type": "VARCHAR"}, # ] + + translate_type("int8", "postgres", "bigquery") # → "INT64" """ from __future__ import annotations @@ -60,3 +62,55 @@ def parse_types( row["type"] = parse_type(row.get(type_field, ""), dialect) results.append(row) return results + + +def translate_type(type_str: str, source_dialect: str, target_dialect: str) -> str: + """Translate a SQL type string from one dialect to another. + + Parses *type_str* using *source_dialect* and re-serializes it in + *target_dialect*, mapping vendor-specific spellings across engines + (e.g. postgres ``int8`` → bigquery ``INT64``, postgres + ``character varying(255)`` → clickhouse ``Nullable(String)``). + + Args: + type_str: Raw type string in the source dialect. + source_dialect: sqlglot dialect to parse with (e.g. "postgres"). + target_dialect: sqlglot dialect to render in (e.g. "bigquery"). + + Returns: + The type string rendered in *target_dialect*. Falls back to the + original string if parsing fails. + """ + if not type_str: + return type_str + try: + parsed = sqlglot.parse_one(type_str, into=DataType, dialect=source_dialect) + except (sqlglot.errors.ParseError, ValueError): + return type_str + try: + return parsed.sql(dialect=target_dialect) + except (sqlglot.errors.ParseError, ValueError): + return type_str + + +def translate_types( + columns: list[dict], + source_dialect: str, + target_dialect: str, + *, + type_field: str = "raw_type", +) -> list[dict]: + """Batch-translate types from *source_dialect* to *target_dialect*. + + Each dict must have a key matching *type_field* (default "raw_type"). + Returns a new list with an added "type" key holding the translated type. + Original dicts are not mutated. + """ + results = [] + for col in columns: + row = dict(col) + row["type"] = translate_type( + row.get(type_field, ""), source_dialect, target_dialect + ) + results.append(row) + return results diff --git a/core/wren/src/wren/utils_cli.py b/core/wren/src/wren/utils_cli.py index d2a4d2b83..c49b9da16 100644 --- a/core/wren/src/wren/utils_cli.py +++ b/core/wren/src/wren/utils_cli.py @@ -51,7 +51,12 @@ def parse_types_cmd( if not path.exists(): typer.echo(f"Error: file not found: {input_file}", err=True) raise typer.Exit(1) - data = json.loads(path.read_text(encoding="utf-8")) + try: + raw = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError) as e: + typer.echo(f"Error: could not read file {input_file}: {e}", err=True) + raise typer.Exit(1) + data = json.loads(raw) else: data = json.load(sys.stdin) except json.JSONDecodeError as e: @@ -60,3 +65,63 @@ def parse_types_cmd( results = parse_types(data, dialect, type_field=type_field) typer.echo(json.dumps(results, indent=2)) + + +@utils_app.command(name="translate-type") +def translate_type_cmd( + type_str: Annotated[str, typer.Option("--type", "-t", help="Raw SQL type string")], + source: Annotated[ + str, + typer.Option("--source", "-s", help="Source SQL dialect (e.g. postgres)"), + ], + target: Annotated[ + str, + typer.Option("--target", help="Target SQL dialect (e.g. bigquery)"), + ], +): + """Translate a single SQL type string from one dialect to another.""" + from wren.type_mapping import translate_type # noqa: PLC0415 + + typer.echo(translate_type(type_str, source, target)) + + +@utils_app.command(name="translate-types") +def translate_types_cmd( + source: Annotated[str, typer.Option("--source", "-s", help="Source SQL dialect")], + target: Annotated[str, typer.Option("--target", help="Target SQL dialect")], + type_field: Annotated[ + str, + typer.Option("--type-field", help="Key name for raw type in input JSON"), + ] = "raw_type", + input_file: Annotated[ + Optional[str], + typer.Option("--input", "-i", help="Input JSON file (default: stdin)"), + ] = None, +): + """Batch-translate types between dialects. Reads/writes JSON. + + Input format: [{"column": "id", "raw_type": "int8"}, ...] + Output format: [{"column": "id", "raw_type": "int8", "type": "INT64"}, ...] + """ + from wren.type_mapping import translate_types # noqa: PLC0415 + + try: + if input_file: + path = Path(input_file) + if not path.exists(): + typer.echo(f"Error: file not found: {input_file}", err=True) + raise typer.Exit(1) + try: + raw = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError) as e: + typer.echo(f"Error: could not read file {input_file}: {e}", err=True) + raise typer.Exit(1) + data = json.loads(raw) + else: + data = json.load(sys.stdin) + except json.JSONDecodeError as e: + typer.echo(f"Error: invalid JSON input: {e}", err=True) + raise typer.Exit(1) + + results = translate_types(data, source, target, type_field=type_field) + typer.echo(json.dumps(results, indent=2)) diff --git a/core/wren/tests/unit/test_type_mapping.py b/core/wren/tests/unit/test_type_mapping.py index 8abdbe414..e8146b42e 100644 --- a/core/wren/tests/unit/test_type_mapping.py +++ b/core/wren/tests/unit/test_type_mapping.py @@ -8,7 +8,12 @@ import sys import pytest -from wren.type_mapping import parse_type, parse_types +from wren.type_mapping import ( + parse_type, + parse_types, + translate_type, + translate_types, +) # ── parse_type unit tests ────────────────────────────────────────────────── @@ -87,6 +92,64 @@ def test_parse_types_empty_list() -> None: assert parse_types([], dialect="postgres") == [] +# ── translate_type cross-dialect tests ────────────────────────── + + +@pytest.mark.parametrize( + "type_str, source, target, expected", + [ + # postgres → bigquery + ("int8", "postgres", "bigquery", "INT64"), + ("TIMESTAMP WITH TIME ZONE", "postgres", "bigquery", "TIMESTAMP"), + # bigquery → postgres round-trip + ("INT64", "bigquery", "postgres", "BIGINT"), + # mysql → snowflake keeps precision/scale + ("DECIMAL(10,2)", "mysql", "snowflake", "DECIMAL(10, 2)"), + # same dialect is an identity-ish normalization + ("int8", "postgres", "postgres", "BIGINT"), + # graceful fallback for unknown types + ("my_custom_type", "postgres", "bigquery", "my_custom_type"), + # empty string passthrough + ("", "postgres", "bigquery", ""), + ], +) +def test_translate_type( + type_str: str, source: str, target: str, expected: str +) -> None: + assert translate_type(type_str, source, target) == expected + + +def test_translate_types_adds_type_field() -> None: + columns = [ + {"column": "id", "raw_type": "int8"}, + {"column": "total", "raw_type": "numeric(10,2)"}, + ] + results = translate_types(columns, "postgres", "bigquery") + + assert len(results) == 2 + assert results[0]["type"] == "INT64" + assert results[1]["type"] == "NUMERIC(10, 2)" + + +def test_translate_types_does_not_mutate_input() -> None: + original = {"column": "id", "raw_type": "int8"} + columns = [original] + translate_types(columns, "postgres", "bigquery") + assert "type" not in original + + +def test_translate_types_custom_type_field() -> None: + columns = [{"col": "x", "data_type": "int8"}] + results = translate_types( + columns, "postgres", "bigquery", type_field="data_type" + ) + assert results[0]["type"] == "INT64" + + +def test_translate_types_empty_list() -> None: + assert translate_types([], "postgres", "bigquery") == [] + + # ── CLI integration tests ───────────────────────────────────────────────── @@ -154,3 +217,89 @@ def test_cli_parse_types_batch() -> None: assert len(data) == 2 assert data[0]["type"] == "BIGINT" assert data[1]["type"] == "VARCHAR" + + +def test_cli_translate_type_single() -> None: + result = _run_wren( + "utils", + "translate-type", + "--type", + "int8", + "--source", + "postgres", + "--target", + "bigquery", + ) + _assert_success(result) + assert result.stdout.strip() == "INT64" + + +def test_cli_translate_type_fallback() -> None: + result = _run_wren( + "utils", + "translate-type", + "--type", + "my_custom_type", + "--source", + "postgres", + "--target", + "bigquery", + ) + _assert_success(result) + assert result.stdout.strip() == "my_custom_type" + + +def test_cli_translate_types_stdin() -> None: + columns = [ + {"column": "id", "raw_type": "int8"}, + {"column": "name", "raw_type": "character varying"}, + ] + result = _run_wren( + "utils", + "translate-types", + "--source", + "postgres", + "--target", + "bigquery", + stdin=json.dumps(columns), + ) + _assert_success(result) + data = json.loads(result.stdout) + assert len(data) == 2 + assert data[0]["type"] == "INT64" + assert data[0]["column"] == "id" + + +def test_cli_translate_types_missing_file() -> None: + result = _run_wren( + "utils", + "translate-types", + "--source", + "postgres", + "--target", + "bigquery", + "--input", + "/nonexistent/does_not_exist.json", + ) + assert result.returncode == 1 + assert "file not found" in result.stderr + assert "Traceback" not in result.stderr + + +def test_cli_translate_types_unreadable_file_is_clean(tmp_path) -> None: + # A directory path is readable-as-path but raises OSError on read_text. + bad = tmp_path / "a_directory" + bad.mkdir() + result = _run_wren( + "utils", + "translate-types", + "--source", + "postgres", + "--target", + "bigquery", + "--input", + str(bad), + ) + assert result.returncode == 1 + assert "could not read file" in result.stderr + assert "Traceback" not in result.stderr