Add with_properties and interpolated_properties

cratedb_sqlparse_py/README.md

@@ -10,20 +10,24 @@
 
 [![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
 ```
 
 ## 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 {<EOF>, ';'}]
-    
+
     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,83 @@ 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.
+
+Properties defined within a `WITH` statement, `statement.metadata.with_properties`.
+
+
+```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 real defined value, marked with a dollar string,  `metadata.interpolated_properties`
+
+```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, `blocks.metadata` will be in `with_properties` and `interpolated_properties` as well.
+
+For values to be picked up they need to start with a dollar `'$'` and be preceded by integers, e.g. `'$1'`, `'$123'` -
+`'$123abc'` would not be valid.
 
 
 ## Development
 
 ### Set up environment
+
 ```shell
 git clone https://github.com/crate/cratedb-sqlparse
 
@@ -125,21 +184,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
 ```

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:

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:

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"