SERVER-61580 Add support for extraData to enums

7 files changed, 210 insertions, 8 deletions

diff --git a/buildscripts/idl/idl/ast.py b/buildscripts/idl/idl/ast.py
index 26238ed53d1..0df99dd093b 100644
--- a/buildscripts/idl/idl/ast.py
+++ b/buildscripts/idl/idl/ast.py
@@ -34,7 +34,7 @@ This is a lossy translation from the IDL Syntax tree as the IDL AST only contain
 the enums and structs that need code generated for them, and just enough information to do that.
 """
 from abc import ABCMeta, abstractmethod
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
 from . import common, errors
 
@@ -341,7 +341,9 @@ class EnumValue(common.SourceLocation):
         # type: (str, int, int) -> None
         """Construct an Enum."""
         self.name = None  # type: str
+        self.description = None  # type: str
         self.value = None  # type: str
+        self.extra_data = None  # type: Dict[str, Any]
 
         super(EnumValue, self).__init__(file_name, line, column)
 
diff --git a/buildscripts/idl/idl/binder.py b/buildscripts/idl/idl/binder.py
index ab50d0bcc35..da38cb176d5 100644
--- a/buildscripts/idl/idl/binder.py
+++ b/buildscripts/idl/idl/binder.py
@@ -1238,7 +1238,9 @@ def _bind_enum(ctxt, idl_enum):
     for enum_value in idl_enum.values:
         ast_enum_value = ast.EnumValue(enum_value.file_name, enum_value.line, enum_value.column)
         ast_enum_value.name = enum_value.name
+        ast_enum_value.description = enum_value.description
         ast_enum_value.value = enum_value.value
+        ast_enum_value.extra_data = enum_value.extra_data
         ast_enum.values.append(ast_enum_value)
 
     values_set = set()  # type: Set[str]
diff --git a/buildscripts/idl/idl/enum_types.py b/buildscripts/idl/idl/enum_types.py
index e9906449f24..4098968ea10 100644
--- a/buildscripts/idl/idl/enum_types.py
+++ b/buildscripts/idl/idl/enum_types.py
@@ -32,8 +32,10 @@ Support the code generation for enums
 """
 
 from abc import ABCMeta, abstractmethod
+import json
 import textwrap
 from typing import cast, List, Optional, Union
+import bson
 
 from . import ast
 from . import common
@@ -90,6 +92,12 @@ class EnumTypeInfoBase(object, metaclass=ABCMeta):
         return "::" + common.qualify_cpp_name(self._enum.cpp_namespace,
                                               self._get_enum_serializer_name())
 
+    def _get_enum_extra_data_name(self):
+        # type: () -> str
+        """Return the name of the get_extra_data function without prefix."""
+        return common.template_args("${enum_name}_get_extra_data",
+                                    enum_name=common.title_case(self._enum.name))
+
     @abstractmethod
     def get_cpp_value_assignment(self, enum_value):
         # type: (ast.EnumValue) -> str
@@ -120,6 +128,66 @@ class EnumTypeInfoBase(object, metaclass=ABCMeta):
         """Generate the serializer function definition."""
         pass
 
+    def _get_populated_extra_values(self):
+        # type: () -> List[Union[syntax.EnumValue,ast.EnumValue]]
+        """Filter the enum values to just those containing extra_data."""
+        return [val for val in self._enum.values if val.extra_data is not None]
+
+    def get_extra_data_declaration(self):
+        # type: () -> Optional[str]
+        """Get the get_extra_data function declaration minus trailing semicolon."""
+        if len(self._get_populated_extra_values()) == 0:
+            return None
+
+        return common.template_args("BSONObj ${function_name}(${enum_name} value)",
+                                    enum_name=self.get_cpp_type_name(),
+                                    function_name=self._get_enum_extra_data_name())
+
+    def gen_extra_data_definition(self, indented_writer):
+        # type: (writer.IndentedTextWriter) -> None
+        """Generate the get_extra_data function definition."""
+
+        extra_values = self._get_populated_extra_values()
+        if len(extra_values) == 0:
+            # No extra data present on this enum.
+            return
+
+        # Generate an anonymous namespace full of BSON constants.
+        #
+        with writer.NamespaceScopeBlock(indented_writer, ['']):
+            for enum_value in extra_values:
+                indented_writer.write_line(
+                    common.template_args('// %s' % json.dumps(enum_value.extra_data)))
+
+                bson_value = ''.join(
+                    [('\\x%02x' % (b)) for b in bson.BSON.encode(enum_value.extra_data)])
+                indented_writer.write_line(
+                    common.template_args(
+                        'const BSONObj ${const_name}("${bson_value}");',
+                        const_name=_get_constant_enum_extra_data_name(self._enum, enum_value),
+                        bson_value=bson_value))
+
+        indented_writer.write_empty_line()
+
+        # Generate implementation of get_extra_data function.
+        #
+        template_params = {
+            'enum_name': self.get_cpp_type_name(),
+            'function_name': self.get_extra_data_declaration(),
+        }
+
+        with writer.TemplateContext(indented_writer, template_params):
+            with writer.IndentedScopedBlock(indented_writer, "${function_name} {", "}"):
+                with writer.IndentedScopedBlock(indented_writer, "switch (value) {", "}"):
+                    for enum_value in extra_values:
+                        indented_writer.write_template(
+                            'case ${enum_name}::%s: return %s;' %
+                            (enum_value.name,
+                             _get_constant_enum_extra_data_name(self._enum, enum_value)))
+                    if len(extra_values) != len(self._enum.values):
+                        # One or more enums does not have associated extra data.
+                        indented_writer.write_line('default: return BSONObj();')
+
 
 class _EnumTypeInt(EnumTypeInfoBase, metaclass=ABCMeta):
     """Type information for integer enumerations."""
@@ -194,6 +262,13 @@ def _get_constant_enum_name(idl_enum, enum_value):
                                 name=enum_value.name)
 
 
+def _get_constant_enum_extra_data_name(idl_enum, enum_value):
+    # type: (Union[syntax.Enum,ast.Enum], Union[syntax.EnumValue,ast.EnumValue]) -> str
+    """Return the C++ name for a string constant of enum extra data value."""
+    return common.template_args('k${enum_name}_${name}_extra_data',
+                                enum_name=common.title_case(idl_enum.name), name=enum_value.name)
+
+
 class _EnumTypeString(EnumTypeInfoBase, metaclass=ABCMeta):
     """Type information for string enumerations."""
 
diff --git a/buildscripts/idl/idl/generator.py b/buildscripts/idl/idl/generator.py
index 30a5f28bf79..d9879e97109 100644
--- a/buildscripts/idl/idl/generator.py
+++ b/buildscripts/idl/idl/generator.py
@@ -678,6 +678,10 @@ class _CppHeaderFileWriter(_CppFileWriterBase):
 
         self._writer.write_line("%s;" % (enum_type_info.get_serializer_declaration()))
 
+        extra_data_decl = enum_type_info.get_extra_data_declaration()
+        if extra_data_decl is not None:
+            self._writer.write_line("%s;" % (extra_data_decl))
+
     def gen_enum_declaration(self, idl_enum):
         # type: (ast.Enum) -> None
         """Generate the declaration for an enum."""
@@ -686,6 +690,8 @@ class _CppHeaderFileWriter(_CppFileWriterBase):
         with self._block('enum class %s : std::int32_t {' % (enum_type_info.get_cpp_type_name()),
                          '};'):
             for enum_value in idl_enum.values:
+                if enum_value.description is not None:
+                    self.gen_description_comment(enum_value.description)
                 self._writer.write_line(
                     common.template_args('${name} ${value},', name=enum_value.name,
                                          value=enum_type_info.get_cpp_value_assignment(enum_value)))
@@ -2205,6 +2211,8 @@ class _CppSourceFileWriter(_CppFileWriterBase):
         enum_type_info.gen_serializer_definition(self._writer)
         self._writer.write_empty_line()
 
+        enum_type_info.gen_extra_data_definition(self._writer)
+
     def _gen_known_fields_declaration(self, struct, name, include_op_msg_implicit):
         # type: (ast.Struct, str, bool) -> None
         """Generate the known fields declaration with specified name."""
diff --git a/buildscripts/idl/idl/parser.py b/buildscripts/idl/idl/parser.py
index e373366f480..d624b4f4473 100644
--- a/buildscripts/idl/idl/parser.py
+++ b/buildscripts/idl/idl/parser.py
@@ -635,6 +635,21 @@ def _parse_generic_reply_field_list_entries(ctxt, node):
     return _parse_field_list_entries(ctxt, node, False)
 
 
+def _parse_arbitrary_value(ctxt, node):
+    # type: (errors.ParserContext, Union[yaml.nodes.MappingNode, yaml.nodes.ScalarNode, yaml.nodes.SequenceNode]) -> Any
+    """Parse a generic YAML type to a Python type."""
+
+    if node.id == 'mapping':
+        return {k.value: _parse_arbitrary_value(ctxt, v) for (k, v) in node.value}
+    elif node.id == 'sequence':
+        return [_parse_arbitrary_value(ctxt, node) for node in node.value]
+    elif ctxt.is_scalar_node(node, 'node'):
+        return node.value
+    else:
+        # Error added to context by is_scalar_node case above
+        return None
+
+
 def _parse_enum_values(ctxt, node):
     # type: (errors.ParserContext, yaml.nodes.MappingNode) -> List[syntax.EnumValue]
     """Parse a values section in an enum in the IDL file."""
@@ -651,13 +666,20 @@ def _parse_enum_values(ctxt, node):
             ctxt.add_duplicate_error(first_node, first_name)
             continue
 
-        # Simple Type
-        if ctxt.is_scalar_node(second_node, first_name):
-            enum_value = syntax.EnumValue(ctxt.file_name, node.start_mark.line,
-                                          node.start_mark.column)
-            enum_value.name = first_name
+        enum_value = syntax.EnumValue(ctxt.file_name, node.start_mark.line, node.start_mark.column)
+        enum_value.name = first_name
+
+        if second_node.id == 'mapping':
+            _generic_parser(
+                ctxt, second_node, first_name, enum_value, {
+                    "description": _RuleDesc('scalar', _RuleDesc.REQUIRED),
+                    "value": _RuleDesc('scalar', _RuleDesc.REQUIRED),
+                    "extra_data": _RuleDesc('mapping', mapping_parser_func=_parse_arbitrary_value),
+                })
+        elif ctxt.is_scalar_node(second_node, first_name):
             enum_value.value = second_node.value
-            enum_values.append(enum_value)
+
+        enum_values.append(enum_value)
 
         field_name_set.add(first_name)
 
@@ -667,7 +689,7 @@ def _parse_enum_values(ctxt, node):
 def _parse_enum(ctxt, spec, name, node):
     # type: (errors.ParserContext, syntax.IDLSpec, str, Union[yaml.nodes.MappingNode, yaml.nodes.ScalarNode, yaml.nodes.SequenceNode]) -> None
     """Parse an enum section in the IDL file."""
-    if not ctxt.is_mapping_node(node, "struct"):
+    if not ctxt.is_mapping_node(node, "enum"):
         return
 
     idl_enum = syntax.Enum(ctxt.file_name, node.start_mark.line, node.start_mark.column)
diff --git a/buildscripts/idl/idl/syntax.py b/buildscripts/idl/idl/syntax.py
index dbc02937631..8558f725c37 100644
--- a/buildscripts/idl/idl/syntax.py
+++ b/buildscripts/idl/idl/syntax.py
@@ -663,7 +663,9 @@ class EnumValue(common.SourceLocation):
         # type: (str, int, int) -> None
         """Construct an Enum."""
         self.name = None  # type: str
+        self.description = None  # type: str
         self.value = None  # type: str
+        self.extra_data = None  # type: Dict[str, Any]
 
         super(EnumValue, self).__init__(file_name, line, column)
 
diff --git a/buildscripts/idl/tests/test_parser.py b/buildscripts/idl/tests/test_parser.py
index 20f67dc0889..1439d58c383 100644
--- a/buildscripts/idl/tests/test_parser.py
+++ b/buildscripts/idl/tests/test_parser.py
@@ -759,6 +759,34 @@ class TestParser(testcase.IDLTestcase):
                     v1: 0
             """))
 
+        # Test extended value
+        self.assert_parse(
+            textwrap.dedent("""
+        enums:
+            foo:
+                description: foo
+                type: foo
+                values:
+                    v1:
+                        description: foo
+                        value: 0
+            """))
+
+        # Test extra_data
+        self.assert_parse(
+            textwrap.dedent("""
+        enums:
+            foo:
+                description: foo
+                type: foo
+                values:
+                    v1:
+                        description: foo
+                        value: 0
+                        extra_data:
+                            bar: baz
+            """))
+
     def test_enum_negative(self):
         # type: () -> None
         """Negative enum test cases."""
@@ -912,6 +940,69 @@ class TestParser(testcase.IDLTestcase):
                     v1: 1
             """), idl.errors.ERROR_ID_DUPLICATE_NODE)
 
+        # Test extra_data invalid type
+        self.assert_parse_fail(
+            textwrap.dedent("""
+        enums:
+            foo:
+                description: foo
+                type: int
+                values:
+                    v1: [ 'foo' ]
+            """), idl.errors.ERROR_ID_IS_NODE_TYPE)
+
+        # Test extended value missing fields (description)
+        self.assert_parse_fail(
+            textwrap.dedent("""
+        enums:
+            foo:
+                description: foo
+                type: int
+                values:
+                    v1:
+                        value: 0
+            """), idl.errors.ERROR_ID_MISSING_REQUIRED_FIELD)
+
+        # Test extended value missing fields (value)
+        self.assert_parse_fail(
+            textwrap.dedent("""
+        enums:
+            foo:
+                description: foo
+                type: int
+                values:
+                    v1:
+                        description: foo
+            """), idl.errors.ERROR_ID_MISSING_REQUIRED_FIELD)
+
+        # Test invalid extra_data (scalar)
+        self.assert_parse_fail(
+            textwrap.dedent("""
+        enums:
+            foo:
+                description: foo
+                type: int
+                values:
+                    v1:
+                        description: foo
+                        value: 0
+                        extra_data: foo
+            """), idl.errors.ERROR_ID_IS_NODE_TYPE)
+
+        # Test invalid extra_data (sequence)
+        self.assert_parse_fail(
+            textwrap.dedent("""
+        enums:
+            foo:
+                description: foo
+                type: int
+                values:
+                    v1:
+                        description: foo
+                        value: 0
+                        extra_data: [ foo ]
+            """), idl.errors.ERROR_ID_IS_NODE_TYPE)
+
     def test_command_positive(self):
         # type: () -> None
         """Positive command test cases."""