From 69fedd41290fcd9d748b8f5dfe7ca3f2a9e71399 Mon Sep 17 00:00:00 2001 From: Ievgenii Shepeliuk Date: Thu, 1 Jun 2023 09:18:40 +0000 Subject: [PATCH] feat: ksql type formatting --- poetry.lock | 59 ++++++++++++++++++++++++++++++++++++++++++- pykli/printer.py | 25 ++++++++++++++++-- pyproject.toml | 3 +++ tests/test_printer.py | 36 ++++++++++++++++++++++++++ 4 files changed, 120 insertions(+), 3 deletions(-) create mode 100644 tests/test_printer.py diff --git a/poetry.lock b/poetry.lock index 8f31121..50aa6c1 100644 --- a/poetry.lock +++ b/poetry.lock @@ -193,6 +193,43 @@ files = [ {file = "idna-3.4.tar.gz", hash = "sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4"}, ] +[[package]] +name = "iniconfig" +version = "2.0.0" +description = "brain-dead simple config-ini parsing" +optional = false +python-versions = ">=3.7" +files = [ + {file = "iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374"}, + {file = "iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3"}, +] + +[[package]] +name = "packaging" +version = "23.1" +description = "Core utilities for Python packages" +optional = false +python-versions = ">=3.7" +files = [ + {file = "packaging-23.1-py3-none-any.whl", hash = "sha256:994793af429502c4ea2ebf6bf664629d07c1a9fe974af92966e4b8d2df7edc61"}, + {file = "packaging-23.1.tar.gz", hash = "sha256:a392980d2b6cffa644431898be54b0045151319d1e7ec34f0cfed48767dd334f"}, +] + +[[package]] +name = "pluggy" +version = "1.0.0" +description = "plugin and hook calling mechanisms for python" +optional = false +python-versions = ">=3.6" +files = [ + {file = "pluggy-1.0.0-py2.py3-none-any.whl", hash = "sha256:74134bbf457f031a36d68416e1509f34bd5ccc019f0bcc952c7b909d06b37bd3"}, + {file = "pluggy-1.0.0.tar.gz", hash = "sha256:4224373bacce55f955a878bf9cfa763c1e360858e330072059e10bad68531159"}, +] + +[package.extras] +dev = ["pre-commit", "tox"] +testing = ["pytest", "pytest-benchmark"] + [[package]] name = "prompt-toolkit" version = "3.0.38" @@ -221,6 +258,26 @@ files = [ [package.extras] plugins = ["importlib-metadata"] +[[package]] +name = "pytest" +version = "7.3.1" +description = "pytest: simple powerful testing with Python" +optional = false +python-versions = ">=3.7" +files = [ + {file = "pytest-7.3.1-py3-none-any.whl", hash = "sha256:3799fa815351fea3a5e96ac7e503a96fa51cc9942c3753cda7651b93c1cfa362"}, + {file = "pytest-7.3.1.tar.gz", hash = "sha256:434afafd78b1d78ed0addf160ad2b77a30d35d4bdf8af234fe621919d9ed15e3"}, +] + +[package.dependencies] +colorama = {version = "*", markers = "sys_platform == \"win32\""} +iniconfig = "*" +packaging = "*" +pluggy = ">=0.12,<2.0" + +[package.extras] +testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "xmlschema"] + [[package]] name = "six" version = "1.16.0" @@ -301,4 +358,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = "^3.11" -content-hash = "91fc9493eee10853bd62eac6d978594b779a0400cb7bebe5c0c27bb2c4aa0942" +content-hash = "05d4b694bc75934ae8ca3cd20460778e96e6ead5068b6b9b116d15f5400707af" diff --git a/pykli/printer.py b/pykli/printer.py index b15091c..93377f0 100644 --- a/pykli/printer.py +++ b/pykli/printer.py @@ -1,6 +1,7 @@ import click import sqlparse from pprint import pformat, pprint +from textwrap import wrap from pygments.token import Token @@ -55,6 +56,26 @@ ), } + +def format_ksql_type(type_def) -> str: + match type_def: + case {"type": "KEY", "schema": {"type": "STRING"}}: + return "VARCHAR (key)" + case {"schema": {"type": "STRING"}}: + return "VARCHAR" + case {"schema": {"type": "STRUCT", "fields": flds}}: + types_str = "\n".join(wrap(', '.join(f['name'] for f in flds), width=70)) + return f"STRUCT<{types_str}>" + case {"type": "KEY", "schema": {"type": tp}}: + return f"{tp} (key)" + case {"type": "HEADER", "headerKey": hdr, "schema": {"type": tp}}: + return f"{tp} (header('{hdr}'))" + case {"schema": {"type": tp}}: + return tp + case _: + return pformat(type_def) + + def print_show(data_type, json): if data_type in KSQL_SHOW_TYPES: data_extractor, headers, row_extractor = KSQL_SHOW_TYPES[data_type] @@ -69,11 +90,11 @@ def print_show(data_type, json): click.secho(f"`show` not implemented for: {data_type}", fg="red") pprint(json) + def print_describe(data): def row_extractor(rows): for r in rows: - yield (r["name"], f"{r['schema']['type']} header('{r['headerKey']}')" if "headerKey" in r else r["schema"]["type"]) - + yield (r["name"], format_ksql_type(r)) ff = format_output(row_extractor(data), DESCRIBE_HEADERS, format_name="psql", preprocessors=(style_output,), header_token=Token.String, odd_row_token=None, even_row_token=None, style=MONOKAI_STYLE, include_default_pygments_style=False) diff --git a/pyproject.toml b/pyproject.toml index 445b84b..1bddaca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,9 @@ sqlparse = "^0.4.4" [tool.poetry.scripts] pykli = "pykli.__main__:main" +[tool.poetry.group.dev.dependencies] +pytest = "^7.3.1" + [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" diff --git a/tests/test_printer.py b/tests/test_printer.py new file mode 100644 index 0000000..a6ff4ad --- /dev/null +++ b/tests/test_printer.py @@ -0,0 +1,36 @@ +from pprint import pformat +from pykli.printer import format_ksql_type + +STRING_COLUMN = { + 'name': 'StrFld', + 'schema': {'fields': None, 'memberSchema': None, 'type': 'STRING'}, +} + +STRING_KEY_COLUMN = { + 'name': 'StrFldKey', 'type': 'KEY', + 'schema': {'fields': None, 'memberSchema': None, 'type': 'STRING'}, +} + +STRUCT_COLUMN = { + 'name': 'Payload', + 'schema': { + 'memberSchema': None, 'type': 'STRUCT', 'fields': [ + {'name': 'Username', 'schema': {'fields': None, 'memberSchema': None, 'type': 'STRING'}}, + {'name': 'IsMobile', 'schema': {'fields': None, 'memberSchema': None, 'type': 'BOOLEAN'}}, + ] + }, +} + +HEADER_COLUMN = { + 'headerKey': 'tenant_id', 'name': '_tenant_id', 'type': 'HEADER', + 'schema': {'fields': None, 'memberSchema': None, 'type': 'BYTES'}, +} + +def test_format_ksql_type(): + assert format_ksql_type({}) == pformat({}) + assert format_ksql_type(None) == pformat(None) + assert format_ksql_type(STRING_COLUMN) == "VARCHAR" + assert format_ksql_type(STRING_KEY_COLUMN) == "VARCHAR (key)" + assert format_ksql_type(STRUCT_COLUMN) == "STRUCT" + assert format_ksql_type(HEADER_COLUMN) == "BYTES (header('tenant_id'))" +