[Python] Handle deprecation and private annotations in codegen (#2773)

## Changes
Handle deprecated and private annotations in codegen.

- Deprecated fields are excluded from codegen
- Private fields are excluded from documentation

There is a trick to transitively mark fields as deprecated (private)
because it's normal for JSON schema to only mark fields as deprecated
(private) and not put the deprecated (private) annotation into
referenced schemas.

## Why
With that, we can continuously check that the Python code is consistent
with the latest JSON schema.

Author: kanterov

.github/workflows/push.yml

@@ -186,6 +186,5 @@ jobs:
           if ! ( git diff --exit-code ); then
             echo "Generated Python code is not up-to-date. Please run 'pushd experimental/python && make codegen' and commit the changes."
 
-            # TODO block PR if this fails once diffs are fixed
-            # exit 1
+            exit 1
           fi

experimental/python/codegen/codegen/generated_dataclass.py

@@ -6,7 +6,7 @@
 
 import codegen.packages as packages
 from codegen.code_builder import CodeBuilder
-from codegen.jsonschema import Property, Schema
+from codegen.jsonschema import Property, Schema, Stage
 from codegen.packages import is_resource
 
 
@@ -96,6 +96,12 @@ class GeneratedField:
     Factory method for creating a default value, used for lists and dicts.
     """
 
+    experimental: bool
+    """
+    If true, the field is experimental and should not be indexed in docs, and
+    be marked as experimental in docstring.
+    """
+
     def __post_init__(self):
         if self.default_factory is not None and self.default is not None:
             raise ValueError("Can't have both default and default_factory", self)
@@ -124,6 +130,7 @@ class GeneratedDataclass:
 
     fields: list[GeneratedField]
     extends: list[GeneratedType]
+    experimental: bool
 
 
 def generate_field(
@@ -147,6 +154,7 @@ def generate_field(
             default=None,
             default_factory="dict",
             create_func_default="None",
+            experimental=prop.stage == Stage.PRIVATE,
         )
     elif field_type.name == "VariableOrList":
         return GeneratedField(
@@ -158,6 +166,7 @@ def generate_field(
             default=None,
             default_factory="list",
             create_func_default="None",
+            experimental=prop.stage == Stage.PRIVATE,
         )
     elif is_required:
         return GeneratedField(
@@ -169,6 +178,7 @@ def generate_field(
             default=None,
             default_factory=None,
             create_func_default=None,
+            experimental=prop.stage == Stage.PRIVATE,
         )
     else:
         return GeneratedField(
@@ -180,6 +190,7 @@ def generate_field(
             default="None",
             default_factory=None,
             create_func_default="None",
+            experimental=prop.stage == Stage.PRIVATE,
         )
 
 
@@ -308,6 +319,7 @@ def generate_dataclass(schema_name: str, schema: Schema) -> GeneratedDataclass:
         description=schema.description,
         fields=fields,
         extends=extends,
+        experimental=schema.stage == Stage.PRIVATE,
     )
 
 
@@ -347,10 +359,10 @@ def _append_dataclass(b: CodeBuilder, generated: GeneratedDataclass):
     b.append(":").newline()
 
     # FIXME should contain class docstring
-    if not generated.description:
+    if not generated.description and not generated.experimental:
         b.indent().append_triple_quote().append_triple_quote().newline().newline()
     else:
-        _append_description(b, generated.description)
+        _append_description(b, generated.description, generated.experimental)
 
 
 def _append_field(b: CodeBuilder, field: GeneratedField):
@@ -428,11 +440,16 @@ def _append_typed_dict(b: CodeBuilder, generated: GeneratedDataclass):
     b.indent().append_triple_quote().append_triple_quote().newline().newline()
 
 
-def _append_description(b: CodeBuilder, description: Optional[str]):
-    if description:
+def _append_description(b: CodeBuilder, description: Optional[str], experimental: bool):
+    if description or experimental:
         b.indent().append_triple_quote().newline()
-        for line in description.split("\n"):
-            b.indent().append(line).newline()
+        if experimental:
+            b.indent().append(":meta private: [EXPERIMENTAL]").newline()
+            if description:
+                b.indent().newline()
+        if description:
+            for line in description.split("\n"):
+                b.indent().append(line).newline()
         b.indent().append_triple_quote().newline()
 
 
@@ -449,7 +466,7 @@ def get_code(generated: GeneratedDataclass) -> str:
 
     for field in generated.fields:
         _append_field(b, field)
-        _append_description(b, field.description)
+        _append_description(b, field.description, field.experimental)
 
         b.newline()
 
@@ -462,7 +479,7 @@ def get_code(generated: GeneratedDataclass) -> str:
 
     for field in generated.fields:
         _append_typed_dict_field(b, field)
-        _append_description(b, field.description)
+        _append_description(b, field.description, field.experimental)
 
         b.newline()
 

experimental/python/codegen/codegen/generated_enum.py

@@ -5,7 +5,7 @@
 import codegen.packages as packages
 from codegen.code_builder import CodeBuilder
 from codegen.generated_dataclass import _append_description
-from codegen.jsonschema import Schema
+from codegen.jsonschema import Schema, Stage
 
 
 @dataclass(kw_only=True)
@@ -14,6 +14,7 @@ class GeneratedEnum:
     package: str
     values: dict[str, str]
     description: Optional[str]
+    experimental: bool
 
 
 def generate_enum(schema_name: str, schema: Schema) -> GeneratedEnum:
@@ -33,6 +34,7 @@ def generate_enum(schema_name: str, schema: Schema) -> GeneratedEnum:
         package=package,
         values=values,
         description=schema.description,
+        experimental=schema.stage == Stage.PRIVATE,
     )
 
 
@@ -46,7 +48,7 @@ def get_code(generated: GeneratedEnum) -> str:
     b.append(f"class {generated.class_name}(Enum):")
     b.newline()
 
-    _append_description(b, generated.description)
+    _append_description(b, generated.description, generated.experimental)
 
     # Example:
     #

experimental/python/codegen/codegen/jsonschema.py

@@ -7,10 +7,16 @@
 import codegen.packages as packages
 
 
+class Stage:
+    PRIVATE = "PRIVATE"
+
+
 @dataclass
 class Property:
     ref: str
     description: Optional[str] = None
+    deprecated: Optional[bool] = None
+    stage: Optional[str] = None
 
 
 class SchemaType(Enum):
@@ -25,6 +31,8 @@ class Schema:
     properties: dict[str, Property] = field(default_factory=dict)
     required: list[str] = field(default_factory=list)
     description: Optional[str] = None
+    deprecated: Optional[bool] = None
+    stage: Optional[str] = None
 
     def __post_init__(self):
         match self.type:
@@ -76,6 +84,11 @@ def _parse_schema(schema: dict) -> Schema:
     schema = _unwrap_variable(schema) or schema
     properties = {}
 
+    def _parse_bool(value) -> Optional[bool]:
+        assert value is None or isinstance(value, bool)
+
+        return value
+
     for k, v in schema.get("properties", {}).items():
         assert v.get("type") is None
         assert v.get("anyOf") is None
@@ -87,6 +100,8 @@ def _parse_schema(schema: dict) -> Schema:
         prop = Property(
             ref=v["$ref"],
             description=v.get("description"),
+            deprecated=_parse_bool(v.get("deprecated")),
+            stage=v.get("x-databricks-preview"),
         )
 
         properties[k] = prop
@@ -102,6 +117,8 @@ def _parse_schema(schema: dict) -> Schema:
         properties=properties,
         required=schema.get("required", []),
         description=schema.get("description"),
+        deprecated=_parse_bool(schema.get("deprecated")),
+        stage=schema.get("x-databricks-preview"),
     )
 
 

experimental/python/codegen/codegen/jsonschema_patch.py

@@ -3,43 +3,15 @@
 from codegen.jsonschema import Schema
 
 REMOVED_FIELDS = {
-    "jobs.RunJobTask": {
-        # all params except job_parameters should be deprecated and should not be supported
-        "jar_params",
-        "notebook_params",
-        "python_params",
-        "spark_submit_params",
-        "python_named_params",
-        "sql_params",
-        "dbt_commands",
-        # except pipeline_params, that is not deprecated
-    },
-    "jobs.TriggerSettings": {
-        # Old table trigger settings name. Deprecated in favor of `table_update`
-        "table",
+    # TODO remove as a follow-up
+    "jobs.Task": {
+        "dashboard_task",
+        "power_bi_task",
     },
     "compute.ClusterSpec": {
         # doesn't work, openapi schema needs to be updated to be enum
         "kind",
     },
-    "jobs.TaskEmailNotifications": {
-        # Deprecated
-        "no_alert_for_skipped_runs",
-    },
-    "jobs.SparkJarTask": {
-        # Deprecated. A value of `false` is no longer supported.
-        "run_as_repl",
-        # Deprecated
-        "jar_uri",
-    },
-    "resources.Pipeline": {
-        # Deprecated
-        "trigger",
-    },
-    "pipelines.PipelineLibrary": {
-        # Deprecated
-        "whl",
-    },
 }
 
 EXTRA_REQUIRED_FIELDS: dict[str, list[str]] = {

experimental/python/codegen/codegen/main.py

@@ -1,4 +1,5 @@
 import argparse
+from dataclasses import replace
 from pathlib import Path
 from textwrap import dedent
 
@@ -19,6 +20,12 @@ def main(output: str):
     schemas = openapi.get_schemas()
     schemas = openapi_patch.add_extra_required_fields(schemas)
     schemas = openapi_patch.remove_unsupported_fields(schemas)
+
+    schemas = _transitively_mark_deprecated_and_private(
+        packages.RESOURCE_TYPES, schemas
+    )
+    # first remove deprecated fields so there are more unused schemas
+    schemas = _remove_deprecated_fields(schemas)
     schemas = _remove_unused_schemas(packages.RESOURCE_TYPES, schemas)
 
     dataclasses, enums = _generate_code(schemas)
@@ -37,6 +44,56 @@ def main(output: str):
         _write_exports(resource, resource_dataclasses, resource_enums, output)
 
 
+def _transitively_mark_deprecated_and_private(
+    roots: list[str],
+    schemas: dict[str, openapi.Schema],
+) -> dict[str, openapi.Schema]:
+    """
+    If schema is only used through deprecated (private) fields, make it as deprecated (private).
+
+    For example, if a field is marked as private, and is excluded from documentation, corresponding
+    dataclasses and enums should be private as well.
+    """
+
+    not_private = _collect_reachable_schemas(roots, schemas, include_private=False)
+    not_deprecated = _collect_reachable_schemas(
+        roots, schemas, include_deprecated=False
+    )
+    new_schemas = {}
+
+    for schema_name, schema in schemas.items():
+        if schema_name not in not_private:
+            schema.stage = openapi.Stage.PRIVATE
+
+        if schema_name not in not_deprecated:
+            schema.deprecated = True
+
+        new_schemas[schema_name] = schema
+
+    return new_schemas
+
+
+def _remove_deprecated_fields(
+    schemas: dict[str, openapi.Schema],
+) -> dict[str, openapi.Schema]:
+    new_schemas = {}
+
+    for name, schema in schemas.items():
+        if schema.type == openapi.SchemaType.OBJECT:
+            new_properties = {}
+            for field_name, field in schema.properties.items():
+                if field.deprecated:
+                    continue
+
+                new_properties[field_name] = field
+
+            new_schemas[name] = replace(schema, properties=new_properties)
+        else:
+            new_schemas[name] = schema
+
+    return new_schemas
+
+
 def _generate_code(
     schemas: dict[str, openapi.Schema],
 ) -> tuple[dict[str, GeneratedDataclass], dict[str, GeneratedEnum]]:
@@ -182,6 +239,12 @@ def _collect_reachable_schemas(
                 if field.ref:
                     name = field.ref.split("/")[-1]
 
+                    if not include_private and field.stage == openapi.Stage.PRIVATE:
+                        continue
+
+                    if not include_deprecated and field.deprecated:
+                        continue
+
                     if name not in reachable:
                         stack.append(name)
 

experimental/python/codegen/codegen_tests/test_generated_dataclass.py

@@ -65,8 +65,10 @@ def test_generate_dataclass():
                 field_name="task_key",
                 param_type_name=variable_or_type(str_type(), is_required=True),
                 type_name=variable_or_type(str_type(), is_required=True),
+                experimental=False,
             ),
         ],
+        experimental=False,
     )
 
 

experimental/python/codegen/codegen_tests/test_generated_enum.py

@@ -17,4 +17,5 @@ def test_generate_enum():
         package="databricks.bundles.jobs._models.my_enum",
         values={"MY_ENUM_VALUE": "myEnumValue"},
         description="enum description",
+        experimental=False,
     )

experimental/python/databricks/bundles/compute/_models/adlsgen2_info.py

@@ -11,7 +11,9 @@
 
 @dataclass(kw_only=True)
 class Adlsgen2Info:
-    """"""
+    """
+    A storage location in Adls Gen2
+    """
 
     destination: VariableOr[str]
     """