Add span support for PARTITION OF fields.

Author: fmguerreiro

src/ast/spans.rs

@@ -34,19 +34,20 @@ use super::{
     ColumnOption, ColumnOptionDef, ConditionalStatementBlock, ConditionalStatements,
     ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex, CreateTable,
     CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr,
-    ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
+    ExprWithAlias, Fetch, ForValues, FromTable, Function, FunctionArg, FunctionArgExpr,
     FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr, HavingBound,
     IfStatement, IlikeSelectItem, IndexColumn, Insert, Interpolate, InterpolateExpr, Join,
     JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView, LimitClause,
     MatchRecognizePattern, Measure, Merge, MergeAction, MergeClause, MergeInsertExpr,
     MergeInsertKind, MergeUpdateExpr, NamedParenthesizedList, NamedWindowDefinition, ObjectName,
     ObjectNamePart, Offset, OnConflict, OnConflictAction, OnInsert, OpenStatement, OrderBy,
-    OrderByExpr, OrderByKind, OutputClause, Partition, PivotValueSource, ProjectionSelect, Query,
-    RaiseStatement, RaiseStatementValue, ReferentialAction, RenameSelectItem, ReplaceSelectElement,
-    ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SqlOption, Statement, Subscript,
-    SymbolDefinition, TableAlias, TableAliasColumnDef, TableConstraint, TableFactor, TableObject,
-    TableOptionsClustered, TableWithJoins, Update, UpdateTableFromKind, Use, Value, Values,
-    ViewColumnDef, WhileStatement, WildcardAdditionalOptions, With, WithFill,
+    OrderByExpr, OrderByKind, OutputClause, Partition, PartitionBoundValue, PivotValueSource,
+    ProjectionSelect, Query, RaiseStatement, RaiseStatementValue, ReferentialAction,
+    RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem,
+    SetExpr, SqlOption, Statement, Subscript, SymbolDefinition, TableAlias, TableAliasColumnDef,
+    TableConstraint, TableFactor, TableObject, TableOptionsClustered, TableWithJoins, Update,
+    UpdateTableFromKind, Use, Value, Values, ViewColumnDef, WhileStatement,
+    WildcardAdditionalOptions, With, WithFill,
 };
 
 /// Given an iterator of spans, return the [Span::union] of all spans.
@@ -554,8 +555,8 @@ impl Spanned for CreateTable {
             cluster_by: _,                      // todo, BigQuery specific
             clustered_by: _,                    // todo, Hive specific
             inherits: _,                        // todo, PostgreSQL specific
-            partition_of: _,                    // todo, PostgreSQL specific
-            for_values: _,                      // todo, PostgreSQL specific
+            partition_of,
+            for_values,
             strict: _,                          // bool
             copy_grants: _,                     // bool
             enable_schema_evolution: _,         // bool
@@ -586,7 +587,9 @@ impl Spanned for CreateTable {
                 .chain(columns.iter().map(|i| i.span()))
                 .chain(constraints.iter().map(|i| i.span()))
                 .chain(query.iter().map(|i| i.span()))
-                .chain(clone.iter().map(|i| i.span())),
+                .chain(clone.iter().map(|i| i.span()))
+                .chain(partition_of.iter().map(|i| i.span()))
+                .chain(for_values.iter().map(|i| i.span())),
         )
     }
 }
@@ -624,6 +627,33 @@ impl Spanned for TableConstraint {
     }
 }
 
+impl Spanned for PartitionBoundValue {
+    fn span(&self) -> Span {
+        match self {
+            PartitionBoundValue::Expr(expr) => expr.span(),
+            // MINVALUE and MAXVALUE are keywords without tracked spans
+            PartitionBoundValue::MinValue => Span::empty(),
+            PartitionBoundValue::MaxValue => Span::empty(),
+        }
+    }
+}
+
+impl Spanned for ForValues {
+    fn span(&self) -> Span {
+        match self {
+            ForValues::In(exprs) => union_spans(exprs.iter().map(|e| e.span())),
+            ForValues::From { from, to } => union_spans(
+                from.iter()
+                    .map(|v| v.span())
+                    .chain(to.iter().map(|v| v.span())),
+            ),
+            // WITH (MODULUS n, REMAINDER r) - u64 values have no spans
+            ForValues::With { .. } => Span::empty(),
+            ForValues::Default => Span::empty(),
+        }
+    }
+}
+
 impl Spanned for CreateIndex {
     fn span(&self) -> Span {
         let CreateIndex {