From b95a0aa45dd936c97989c401f0ee915f17543915 Mon Sep 17 00:00:00 2001 From: surister Date: Tue, 11 Jun 2024 19:02:44 +0200 Subject: [PATCH] Add `with_properties` and `interpolated_properties` --- cratedb_sqlparse_py/README.md | 85 ++++++++++++++++--- .../cratedb_sqlparse/AstBuilder.py | 20 +++++ .../cratedb_sqlparse/parser.py | 2 + cratedb_sqlparse_py/tests/test_enricher.py | 29 +++++++ 4 files changed, 126 insertions(+), 10 deletions(-) diff --git a/cratedb_sqlparse_py/README.md b/cratedb_sqlparse_py/README.md index ee99acd..a56a5dd 100644 --- a/cratedb_sqlparse_py/README.md +++ b/cratedb_sqlparse_py/README.md @@ -10,13 +10,16 @@ [![Status](https://img.shields.io/pypi/status/cratedb-sqlparse.svg)](https://pypi.org/project/cratedb-sqlparse/) -This package provides utilities to validate and split SQL statements specifically designed for CrateDB. +This package provides utilities to validate and split SQL statements specifically designed for +CrateDB. -It is built upon CrateDB's antlr4 grammar, ensuring accurate parsing tailored to CrateDB's SQL dialect. +It is built upon CrateDB's antlr4 grammar, ensuring accurate parsing tailored to CrateDB's SQL +dialect. It draws inspiration from `sqlparse`. ## Installation. + ```shell pip install cratedb-sqlparse ``` @@ -24,6 +27,7 @@ pip install cratedb-sqlparse ## Usage. ### Simple example + ```python from cratedb_sqlparse import sqlparse @@ -49,7 +53,9 @@ print(select_query.tree) ``` ### Exceptions and errors. + By default exceptions are stored in `statement.exception` + ```python from cratedb_sqlparse import sqlparse @@ -64,19 +70,17 @@ stmt = statements[0] if stmt.exception: print(stmt.exception.error_message) # InputMismatchException[line 2:31 mismatched input 'HERE' expecting {, ';'}] - + print(stmt.exception.original_query_with_error_marked) # SELECT COUNT(*) FROM doc.tbl f HERE f.id = 1; # ^^^^ # # INSERT INTO doc.tbl VALUES (1, 23, 4); - + print(stmt.exception.offending_token.text) # HERE - ``` - In some situations, you might want sqlparse to raise an exception. You can set `raise_exception` to `True` @@ -87,28 +91,85 @@ from cratedb_sqlparse import sqlparse sqlparse('SELECT COUNT(*) FROM doc.tbl f WHERE .id = 1;', raise_exception=True) # cratedb_sqlparse.parser.ParsingException: NoViableAltException[line 1:37 no viable alternative at input 'SELECT COUNT(*) FROM doc.tbl f WHERE .'] -``` +``` Catch the exception: -```python +```python from cratedb_sqlparse import sqlparse, ParsingException try: t = sqlparse('SELECT COUNT(*) FROM doc.tbl f WHERE .id = 1;', raise_exception=True)[0] except ParsingException: print('Catched!') - ``` Note: -It will only raise the first exception if finds, even if you pass in several statements. +It will only raise the first exception it finds, even if you pass in several statements. + +### Query metadata. + +Query metadata can be read with `statement.metadata` + +```python +from cratedb_sqlparse import sqlparse + +stmt = sqlparse("SELECT A, B FROM doc.tbl12") + +print(stmt.metadata) +# Metadata(schema='doc', table_name='tbl12', interpolated_properties={}, with_properties={}) +``` + +#### Query properties. + +Query properties can be read with `statement.metadata.with_properties`. + +Properties are only detected in statements with `WITH`. + +```python +from cratedb_sqlparse import sqlparse + +stmt = sqlparse(""" + CREATE TABLE doc.tbl12 (A TEXT) WITH ( + "allocation.max_retries" = 5, + "blocks.metadata" = false + ); +""")[0] + +print(stmt.metadata) +# Metadata(schema='doc', table_name='tbl12', interpolated_properties={}, with_properties={'allocation.max_retries': '5', 'blocks.metadata': 'false'}) +``` + +#### Interpolated properties. + +Interpolated properties are properties without a defined value, they + +```python +from cratedb_sqlparse import sqlparse + +stmt = sqlparse(""" + CREATE TABLE doc.tbl12 (A TEXT) WITH ( + "allocation.max_retries" = 5, + "blocks.metadata" = $1 +); +""")[0] + +print(stmt.metadata) +# Metadata(schema='doc', table_name='tbl12', interpolated_properties={'blocks.metadata': '$1'}, with_properties={'allocation.max_retries': '5', 'blocks.metadata': '$1'}) +``` + +In this case, both `blocks.metadata` will be in `with_properties` and `interpolated_properties`. +For values to be picked up they need to start with a dollar `'$'` and be numeric, e.g. `'$1'`, `'$123'` - +`'$123abc'` would not be valid. + + ## Development ### Set up environment + ```shell git clone https://github.com/crate/cratedb-sqlparse @@ -125,21 +186,25 @@ Everytime you open a shell again you would need to run `source .venv/bin/activat to use `poe` commands. ### Run lint and tests with coverage. + ```shell poe check ``` ### Run only tests + ```shell poe test ``` ### Run a specific test. + ```shell poe test -k test_sqlparse_collects_exceptions_2 ``` ### Run linter + ```shell poe lint ``` \ No newline at end of file diff --git a/cratedb_sqlparse_py/cratedb_sqlparse/AstBuilder.py b/cratedb_sqlparse_py/cratedb_sqlparse/AstBuilder.py index dd0f796..1203cb3 100644 --- a/cratedb_sqlparse_py/cratedb_sqlparse/AstBuilder.py +++ b/cratedb_sqlparse_py/cratedb_sqlparse/AstBuilder.py @@ -43,6 +43,26 @@ def visitTableName(self, ctx: SqlBaseParser.TableNameContext): self.stmt.metadata.table_name = name self.stmt.metadata.schema = schema + def visitGenericProperties(self, ctx: SqlBaseParser.GenericPropertiesContext): + node_properties = ctx.genericProperty() + + properties = {} + interpolated_properties = {} + + for property in node_properties: + key = self.get_text(property.ident()) + value = self.get_text(property.expr()) + + properties[key] = value + + if value[0] == "$": + # It might be a interpolated value, e.g. '$1' + if value[1:].isdigit(): + interpolated_properties[key] = value + + self.stmt.metadata.with_properties = properties + self.stmt.metadata.interpolated_properties = interpolated_properties + def get_text(self, node) -> t.Optional[str]: """Gets the text representation of the node or None if it doesn't have one""" if node: diff --git a/cratedb_sqlparse_py/cratedb_sqlparse/parser.py b/cratedb_sqlparse_py/cratedb_sqlparse/parser.py index 20e5ca8..a56bf72 100644 --- a/cratedb_sqlparse_py/cratedb_sqlparse/parser.py +++ b/cratedb_sqlparse_py/cratedb_sqlparse/parser.py @@ -134,6 +134,8 @@ class Metadata: schema: str = None table_name: str = None + interpolated_properties: dict = dataclasses.field(default_factory=dict) + with_properties: dict = dataclasses.field(default_factory=dict) class Statement: diff --git a/cratedb_sqlparse_py/tests/test_enricher.py b/cratedb_sqlparse_py/tests/test_enricher.py index 8d02d4c..8ef145b 100644 --- a/cratedb_sqlparse_py/tests/test_enricher.py +++ b/cratedb_sqlparse_py/tests/test_enricher.py @@ -44,3 +44,32 @@ def test_table_name_statements(): assert stmts[3].metadata.schema is None assert stmts[3].metadata.table_name == "tbl1" + + +def test_table_with_properties(): + from cratedb_sqlparse import sqlparse + + query = "CREATE TABLE tbl (A TEXT) WITH ('key' = 'val', 'key2' = 2, 'key3' = true)" + + stmt = sqlparse(query)[0] + keys = ["key", "key2"] + + assert all(x in stmt.metadata.with_properties for x in keys) + assert stmt.metadata.with_properties["key"] == "val" + assert stmt.metadata.with_properties["key2"] == "2" + assert stmt.metadata.with_properties["key3"] == "true" + + +def test_with_with_interpolated_properties(): + from cratedb_sqlparse import sqlparse + + query = "CREATE TABLE tbl (A TEXT) WITH ('key' = $1, 'key2' = '$2')" + + stmt = sqlparse(query)[0] + keys = ["key", "key2"] + + # Has all the keys. + assert all(x in stmt.metadata.interpolated_properties for x in keys) + assert all(x in stmt.metadata.with_properties for x in keys) + assert stmt.metadata.with_properties["key"] == "$1" + assert stmt.metadata.with_properties["key2"] == "$2"