Merge pull request #10130 from Mytherin/querylocation

For parsed expressions and table references - serialize query location, and obtain query location for more properties in transformer
duckdb · Jan 4, 2024 · 6b4f65a · 6b4f65a
2 parents 28543c9 + ddeab6d
commit 6b4f65a
Show file tree

Hide file tree

Showing 11 changed files with 63 additions and 28 deletions.
diff --git a/src/include/duckdb/storage/serialization/parsed_expression.json b/src/include/duckdb/storage/serialization/parsed_expression.json
@@ -24,6 +24,12 @@
         "id": 102,
         "name": "alias",
         "type": "string"
+      },
+      {
+        "id": 103,
+        "name": "query_location",
+        "type": "idx_t",
+        "default": "idx_t(DConstants::INVALID_INDEX)"
       }
     ]
   },

diff --git a/src/include/duckdb/storage/serialization/tableref.json b/src/include/duckdb/storage/serialization/tableref.json
@@ -20,6 +20,12 @@
         "id": 102,
         "name": "sample",
         "type": "SampleOptions*"
+      },
+      {
+        "id": 103,
+        "name": "query_location",
+        "type": "idx_t",
+        "default": "idx_t(DConstants::INVALID_INDEX)"
       }
     ]
   },

diff --git a/src/parser/transform/expression/transform_bool_expr.cpp b/src/parser/transform/expression/transform_bool_expr.cpp
@@ -46,6 +46,7 @@ unique_ptr<ParsedExpression> Transformer::TransformBoolExpr(duckdb_libpgquery::P
 		}
 		}
 	}
+	result->query_location = root.location;
 	return result;
 }
 

diff --git a/src/parser/transform/expression/transform_boolean_test.cpp b/src/parser/transform/expression/transform_boolean_test.cpp
@@ -6,31 +6,46 @@
 
 namespace duckdb {
 
+static unique_ptr<ParsedExpression> TransformBooleanTestInternal(unique_ptr<ParsedExpression> argument,
+                                                                 ExpressionType comparison_type, bool comparison_value,
+                                                                 idx_t query_location) {
+	auto bool_value = make_uniq<ConstantExpression>(Value::BOOLEAN(comparison_value));
+	bool_value->query_location = query_location;
+	// we cast the argument to bool to remove ambiguity wrt function binding on the comparision
+	auto cast_argument = make_uniq<CastExpression>(LogicalType::BOOLEAN, std::move(argument));
+
+	auto result = make_uniq<ComparisonExpression>(comparison_type, std::move(cast_argument), std::move(bool_value));
+	result->query_location = query_location;
+	return std::move(result);
+}
+
+static unique_ptr<ParsedExpression> TransformBooleanTestIsNull(unique_ptr<ParsedExpression> argument,
+                                                               ExpressionType operator_type, idx_t query_location) {
+	auto result = make_uniq<OperatorExpression>(operator_type, std::move(argument));
+	result->query_location = query_location;
+	return std::move(result);
+}
+
 unique_ptr<ParsedExpression> Transformer::TransformBooleanTest(duckdb_libpgquery::PGBooleanTest &node) {
 	auto argument = TransformExpression(PGPointerCast<duckdb_libpgquery::PGNode>(node.arg));
 
-	auto expr_true = make_uniq<ConstantExpression>(Value::BOOLEAN(true));
-	auto expr_false = make_uniq<ConstantExpression>(Value::BOOLEAN(false));
-	// we cast the argument to bool to remove ambiguity wrt function binding on the comparision
-	auto cast_argument = make_uniq<CastExpression>(LogicalType::BOOLEAN, argument->Copy());
-
 	switch (node.booltesttype) {
 	case duckdb_libpgquery::PGBoolTestType::PG_IS_TRUE:
-		return make_uniq<ComparisonExpression>(ExpressionType::COMPARE_NOT_DISTINCT_FROM, std::move(cast_argument),
-		                                       std::move(expr_true));
+		return TransformBooleanTestInternal(std::move(argument), ExpressionType::COMPARE_NOT_DISTINCT_FROM, true,
+		                                    node.location);
 	case duckdb_libpgquery::PGBoolTestType::IS_NOT_TRUE:
-		return make_uniq<ComparisonExpression>(ExpressionType::COMPARE_DISTINCT_FROM, std::move(cast_argument),
-		                                       std::move(expr_true));
+		return TransformBooleanTestInternal(std::move(argument), ExpressionType::COMPARE_DISTINCT_FROM, true,
+		                                    node.location);
 	case duckdb_libpgquery::PGBoolTestType::IS_FALSE:
-		return make_uniq<ComparisonExpression>(ExpressionType::COMPARE_NOT_DISTINCT_FROM, std::move(cast_argument),
-		                                       std::move(expr_false));
+		return TransformBooleanTestInternal(std::move(argument), ExpressionType::COMPARE_NOT_DISTINCT_FROM, false,
+		                                    node.location);
 	case duckdb_libpgquery::PGBoolTestType::IS_NOT_FALSE:
-		return make_uniq<ComparisonExpression>(ExpressionType::COMPARE_DISTINCT_FROM, std::move(cast_argument),
-		                                       std::move(expr_false));
+		return TransformBooleanTestInternal(std::move(argument), ExpressionType::COMPARE_DISTINCT_FROM, false,
+		                                    node.location);
 	case duckdb_libpgquery::PGBoolTestType::IS_UNKNOWN: // IS NULL
-		return make_uniq<OperatorExpression>(ExpressionType::OPERATOR_IS_NULL, std::move(argument));
+		return TransformBooleanTestIsNull(std::move(argument), ExpressionType::OPERATOR_IS_NULL, node.location);
 	case duckdb_libpgquery::PGBoolTestType::IS_NOT_UNKNOWN: // IS NOT NULL
-		return make_uniq<OperatorExpression>(ExpressionType::OPERATOR_IS_NOT_NULL, std::move(argument));
+		return TransformBooleanTestIsNull(std::move(argument), ExpressionType::OPERATOR_IS_NOT_NULL, node.location);
 	default:
 		throw NotImplementedException("Unknown boolean test type %d", node.booltesttype);
 	}

diff --git a/src/parser/transform/expression/transform_case.cpp b/src/parser/transform/expression/transform_case.cpp
@@ -29,6 +29,7 @@ unique_ptr<ParsedExpression> Transformer::TransformCase(duckdb_libpgquery::PGCas
 	} else {
 		case_node->else_expr = make_uniq<ConstantExpression>(Value(LogicalType::SQLNULL));
 	}
+	case_node->query_location = root.location;
 	return std::move(case_node);
 }
 

diff --git a/src/parser/transform/expression/transform_constant.cpp b/src/parser/transform/expression/transform_constant.cpp
@@ -76,7 +76,9 @@ unique_ptr<ConstantExpression> Transformer::TransformValue(duckdb_libpgquery::PG
 }
 
 unique_ptr<ParsedExpression> Transformer::TransformConstant(duckdb_libpgquery::PGAConst &c) {
-	return TransformValue(c.val);
+	auto constant = TransformValue(c.val);
+	constant->query_location = c.location;
+	return std::move(constant);
 }
 
 bool Transformer::ConstructConstantFromExpression(const ParsedExpression &expr, Value &value) {

diff --git a/src/parser/transform/expression/transform_is_null.cpp b/src/parser/transform/expression/transform_is_null.cpp
@@ -13,7 +13,9 @@ unique_ptr<ParsedExpression> Transformer::TransformNullTest(duckdb_libpgquery::P
 	                               ? ExpressionType::OPERATOR_IS_NULL
 	                               : ExpressionType::OPERATOR_IS_NOT_NULL;
 
-	return unique_ptr<ParsedExpression>(new OperatorExpression(expr_type, std::move(arg)));
+	auto result = make_uniq<OperatorExpression>(expr_type, std::move(arg));
+	result->query_location = root.location;
+	return std::move(result);
 }
 
 } // namespace duckdb
diff --git a/src/parser/transform/expression/transform_lambda.cpp b/src/parser/transform/expression/transform_lambda.cpp
@@ -12,7 +12,9 @@ unique_ptr<ParsedExpression> Transformer::TransformLambda(duckdb_libpgquery::PGL
 	auto rhs = TransformExpression(node.rhs);
 	D_ASSERT(lhs);
 	D_ASSERT(rhs);
-	return make_uniq<LambdaExpression>(std::move(lhs), std::move(rhs));
+	auto result = make_uniq<LambdaExpression>(std::move(lhs), std::move(rhs));
+	result->query_location = node.location;
+	return std::move(result);
 }
 
 } // namespace duckdb
diff --git a/src/storage/serialization/serialize_parsed_expression.cpp b/src/storage/serialization/serialize_parsed_expression.cpp
@@ -13,12 +13,14 @@ void ParsedExpression::Serialize(Serializer &serializer) const {
 	serializer.WriteProperty<ExpressionClass>(100, "class", expression_class);
 	serializer.WriteProperty<ExpressionType>(101, "type", type);
 	serializer.WritePropertyWithDefault<string>(102, "alias", alias);
+	serializer.WritePropertyWithDefault<idx_t>(103, "query_location", query_location, idx_t(DConstants::INVALID_INDEX));
 }
 
 unique_ptr<ParsedExpression> ParsedExpression::Deserialize(Deserializer &deserializer) {
 	auto expression_class = deserializer.ReadProperty<ExpressionClass>(100, "class");
 	auto type = deserializer.ReadProperty<ExpressionType>(101, "type");
 	auto alias = deserializer.ReadPropertyWithDefault<string>(102, "alias");
+	auto query_location = deserializer.ReadPropertyWithDefault<idx_t>(103, "query_location", idx_t(DConstants::INVALID_INDEX));
 	deserializer.Set<ExpressionType>(type);
 	unique_ptr<ParsedExpression> result;
 	switch (expression_class) {
@@ -81,6 +83,7 @@ unique_ptr<ParsedExpression> ParsedExpression::Deserialize(Deserializer &deseria
 	}
 	deserializer.Unset<ExpressionType>();
 	result->alias = std::move(alias);
+	result->query_location = query_location;
 	return result;
 }
 

diff --git a/src/storage/serialization/serialize_tableref.cpp b/src/storage/serialization/serialize_tableref.cpp
@@ -13,12 +13,14 @@ void TableRef::Serialize(Serializer &serializer) const {
 	serializer.WriteProperty<TableReferenceType>(100, "type", type);
 	serializer.WritePropertyWithDefault<string>(101, "alias", alias);
 	serializer.WritePropertyWithDefault<unique_ptr<SampleOptions>>(102, "sample", sample);
+	serializer.WritePropertyWithDefault<idx_t>(103, "query_location", query_location, idx_t(DConstants::INVALID_INDEX));
 }
 
 unique_ptr<TableRef> TableRef::Deserialize(Deserializer &deserializer) {
 	auto type = deserializer.ReadProperty<TableReferenceType>(100, "type");
 	auto alias = deserializer.ReadPropertyWithDefault<string>(101, "alias");
 	auto sample = deserializer.ReadPropertyWithDefault<unique_ptr<SampleOptions>>(102, "sample");
+	auto query_location = deserializer.ReadPropertyWithDefault<idx_t>(103, "query_location", idx_t(DConstants::INVALID_INDEX));
 	unique_ptr<TableRef> result;
 	switch (type) {
 	case TableReferenceType::BASE_TABLE:
@@ -47,6 +49,7 @@ unique_ptr<TableRef> TableRef::Deserialize(Deserializer &deserializer) {
 	}
 	result->alias = std::move(alias);
 	result->sample = std::move(sample);
+	result->query_location = query_location;
 	return result;
 }
 

diff --git a/test/sql/json/test_json_serialize_sql.test b/test/sql/json/test_json_serialize_sql.test
@@ -4,22 +4,16 @@
 require json
 
 # Example with simple query
-query I
+statement ok
 SELECT json_serialize_sql('SELECT 1 + 2 FROM tbl1');
-----
-{"error":false,"statements":[{"node":{"type":"SELECT_NODE","modifiers":[],"cte_map":{"map":[]},"select_list":[{"class":"FUNCTION","type":"FUNCTION","alias":"","function_name":"+","schema":"","children":[{"class":"CONSTANT","type":"VALUE_CONSTANT","alias":"","value":{"type":{"id":"INTEGER","type_info":null},"is_null":false,"value":1}},{"class":"CONSTANT","type":"VALUE_CONSTANT","alias":"","value":{"type":{"id":"INTEGER","type_info":null},"is_null":false,"value":2}}],"filter":null,"order_bys":{"type":"ORDER_MODIFIER","orders":[]},"distinct":false,"is_operator":true,"export_state":false,"catalog":""}],"from_table":{"type":"BASE_TABLE","alias":"","sample":null,"schema_name":"","table_name":"tbl1","column_name_alias":[],"catalog_name":""},"where_clause":null,"group_expressions":[],"group_sets":[],"aggregate_handling":"STANDARD_HANDLING","having":null,"sample":null,"qualify":null}}]}
 
 # Example with skip_null and skip_empty
-query I
+statement ok
 SELECT json_serialize_sql('SELECT *, 1 + 2 FROM tbl1', skip_null := true, skip_empty := true);
-----
-{"error":false,"statements":[{"node":{"type":"SELECT_NODE","select_list":[{"class":"STAR","type":"STAR","columns":false},{"class":"FUNCTION","type":"FUNCTION","function_name":"+","children":[{"class":"CONSTANT","type":"VALUE_CONSTANT","value":{"type":{"id":"INTEGER"},"is_null":false,"value":1}},{"class":"CONSTANT","type":"VALUE_CONSTANT","value":{"type":{"id":"INTEGER"},"is_null":false,"value":2}}],"order_bys":{"type":"ORDER_MODIFIER"},"distinct":false,"is_operator":true,"export_state":false}],"from_table":{"type":"BASE_TABLE","table_name":"tbl1"},"aggregate_handling":"STANDARD_HANDLING"}}]}
 
 # Example with subquery
-query I
+statement ok
 SELECT json_serialize_sql('SELECT * FROM (SELECT 1 + 2)', skip_null := true, skip_empty := true);
-----
-{"error":false,"statements":[{"node":{"type":"SELECT_NODE","select_list":[{"class":"STAR","type":"STAR","columns":false}],"from_table":{"type":"SUBQUERY","subquery":{"node":{"type":"SELECT_NODE","select_list":[{"class":"FUNCTION","type":"FUNCTION","function_name":"+","children":[{"class":"CONSTANT","type":"VALUE_CONSTANT","value":{"type":{"id":"INTEGER"},"is_null":false,"value":1}},{"class":"CONSTANT","type":"VALUE_CONSTANT","value":{"type":{"id":"INTEGER"},"is_null":false,"value":2}}],"order_bys":{"type":"ORDER_MODIFIER"},"distinct":false,"is_operator":true,"export_state":false}],"from_table":{"type":"EMPTY"},"aggregate_handling":"STANDARD_HANDLING"}}},"aggregate_handling":"STANDARD_HANDLING"}}]}
 
 # Example with syntax error
 query I
@@ -88,7 +82,7 @@ PRAGMA json_execute_serialized_sql(
 statement error
 SELECT * FROM json_execute_serialized_sql(json_serialize_sql('SELECT * FROM tbl2', skip_null := true, skip_empty := true));
 ----
-Parser Error: Expected but did not find property 'cte_map' in json object: '{"type":"SELECT_NODE","select_list":[{"class":"STAR","type":"STAR","columns":false}],"from_table":{"type":"BASE_TABLE","table_name":"tbl2"},"aggregate_handling":"STANDARD_HANDLING"}'
+Parser Error: Expected but did not find property 'cte_map' in json object
 
 
 # Test execute json serialized sql with multiple nested type tags