Add WITH RecordVisibilityContext

Update grammar to support `WITH RecordVisibilityContext` in SOQL
queries.

Author: cwarden

formatter/soql_test.go

@@ -274,6 +274,45 @@ func TestSOQL(t *testing.T) {
 				`[SELECT Name FROM Account WHERE Id = :accountId WITH SYSTEM_MODE]`,
 				`[SELECT Name FROM Account WHERE Id = :accountId WITH SYSTEM_MODE]`,
 			},
+			{
+				`[SELECT Id, RecordVisibility.VisibilityAttribute FROM Account WHERE Id = 'xxx' WITH RecordVisibilityContext (maxDescriptorPerRecord=100, supportsDomains=true, supportsDelegates=true)]`,
+				`[
+	SELECT
+		Id,
+		RecordVisibility.VisibilityAttribute
+	FROM
+		Account
+	WHERE
+		Id = 'xxx'
+	WITH RecordVisibilityContext (maxDescriptorPerRecord=100, supportsDomains=true, supportsDelegates=true)
+]`,
+			},
+			{
+				`[SELECT recordId, VisibilityAttribute FROM RecordVisibility WHERE recordId = 'xxx' WITH RecordVisibilityContext (maxDescriptorPerRecord=100, supportsDomains=true, supportsDelegates=true)]`,
+				`[
+	SELECT
+		recordId,
+		VisibilityAttribute
+	FROM
+		RecordVisibility
+	WHERE
+		recordId = 'xxx'
+	WITH RecordVisibilityContext (maxDescriptorPerRecord=100, supportsDomains=true, supportsDelegates=true)
+]`,
+			},
+			{
+				`[SELECT Id, RecordVisibility.VisibilityAttribute FROM Account WHERE Id = 'xxx' WITH RecordVisibilityContext (maxDescriptorPerRecord=100)]`,
+				`[
+	SELECT
+		Id,
+		RecordVisibility.VisibilityAttribute
+	FROM
+		Account
+	WHERE
+		Id = 'xxx'
+	WITH RecordVisibilityContext (maxDescriptorPerRecord=100)
+]`,
+			},
 			{
 				`[SELECT StageName, COUNT(Id) cnt FROM opportunity GROUP BY StageName HAVING COUNT(Id) > 1]`,
 				`[

formatter/visitors.go

@@ -1009,9 +1009,31 @@ func (v *FormatVisitor) VisitWithClause(ctx *parser.WithClauseContext) interface
 	if ctx.SYSTEM_MODE() != nil {
 		return "WITH SYSTEM_MODE"
 	}
+	if ctx.RECORDVISIBILITYCONTEXT() != nil {
+		params := make([]string, 0, len(ctx.AllRecordVisibilityContextParam()))
+		for _, param := range ctx.AllRecordVisibilityContextParam() {
+			params = append(params, v.visitRule(param).(string))
+		}
+		return fmt.Sprintf("WITH RecordVisibilityContext (%s)", strings.Join(params, ", "))
+	}
 	return "WITH"
 }
 
+func (v *FormatVisitor) VisitRecordVisibilityContextParam(ctx *parser.RecordVisibilityContextParamContext) interface{} {
+	if intValue := ctx.IntegerLiteral(); intValue != nil {
+		return fmt.Sprintf("maxDescriptorPerRecord=%s", intValue.GetText())
+	}
+	if boolValue := ctx.BooleanLiteral(); boolValue != nil {
+		switch {
+		case ctx.SUPPORTSDOMAINS() != nil:
+			return fmt.Sprintf("supportsDomains=%s", boolValue.GetText())
+		case ctx.SUPPORTSDELEGATES() != nil:
+			return fmt.Sprintf("supportsDelegates=%s", boolValue.GetText())
+		}
+	}
+	return ctx.GetText()
+}
+
 func (v *FormatVisitor) VisitWhereClause(ctx *parser.WhereClauseContext) interface{} {
 	sep := " "
 	indent := 0

grammar/ApexLexer.g4

@@ -155,6 +155,10 @@ ABOVE             : A B O V E;
 BELOW             : B E L O W;
 ABOVE_OR_BELOW    : A B O V E '_' O R '_' B E L O W;
 SECURITY_ENFORCED : S E C U R I T Y '_' E N F O R C E D;
+RECORDVISIBILITYCONTEXT : R E C O R D V I S I B I L I T Y C O N T E X T;
+MAXDESCRIPTORPERRECORD : M A X D E S C R I P T O R P E R R E C O R D;
+SUPPORTSDOMAINS : S U P P O R T S D O M A I N S;
+SUPPORTSDELEGATES : S U P P O R T S D E L E G A T E S;
 REFERENCE         : R E F E R E N C E;
 CUBE              : C U B E;
 FORMAT            : F O R M A T;

grammar/ApexParser.g4

@@ -719,8 +719,15 @@ withClause
     | WITH SECURITY_ENFORCED
     | WITH USER_MODE
     | WITH SYSTEM_MODE
+    | WITH RECORDVISIBILITYCONTEXT LPAREN recordVisibilityContextParam (COMMA recordVisibilityContextParam)* RPAREN
     | WITH logicalExpression;
 
+recordVisibilityContextParam
+    : MAXDESCRIPTORPERRECORD ASSIGN IntegerLiteral
+    | SUPPORTSDOMAINS ASSIGN BooleanLiteral
+    | SUPPORTSDELEGATES ASSIGN BooleanLiteral
+    ;
+
 filteringExpression
     : dataCategorySelection (AND dataCategorySelection)*;
 
@@ -1019,6 +1026,10 @@ id
     | BELOW
     | ABOVE_OR_BELOW
     | SECURITY_ENFORCED
+    | RECORDVISIBILITYCONTEXT
+    | MAXDESCRIPTORPERRECORD
+    | SUPPORTSDOMAINS
+    | SUPPORTSDELEGATES
     | USER_MODE
     | SYSTEM_MODE
     | REFERENCE
@@ -1233,6 +1244,10 @@ anyId
     | BELOW
     | ABOVE_OR_BELOW
     | SECURITY_ENFORCED
+    | RECORDVISIBILITYCONTEXT
+    | MAXDESCRIPTORPERRECORD
+    | SUPPORTSDOMAINS
+    | SUPPORTSDELEGATES
     | USER_MODE
     | SYSTEM_MODE
     | REFERENCE