atlas-init 0.1.0__py3-none-any.whl → 0.1.4__py3-none-any.whl

atlas_init/cli_tf/schema_v3.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+from enum import Enum, StrEnum
+from typing import TypeAlias
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from atlas_init.humps import pascalize
+
+
+def lowercase_no_snake(name: str) -> str:
+    return name.lower().replace("_", "")
+
+
+def default_config_dict() -> ConfigDict:
+    return ConfigDict(populate_by_name=True, alias_generator=lowercase_no_snake)
+
+
+class BaseModelLocal(BaseModel):
+    model_config = default_config_dict()
+
+
+class ElemType(int, Enum):
+    BOOL = 0
+    FLOAT64 = 1
+    INT64 = 2
+    NUMBER = 3
+    STRING = 4
+    UNKNOWN = 5
+
+
+class CustomDefault(BaseModelLocal):
+    definition: str
+    imports: list[str]
+
+
+SnakeCaseString = str
+
+
+class ComputedOptionalRequired(StrEnum):
+    computed = "computed"
+    computed_optional = "computed_optional"
+    optional = "optional"
+    required = "required"
+    unset = ""
+
+
+class BoolAttribute(BaseModelLocal):
+    default: bool | None = None
+
+
+class Float64Attribute(BaseModelLocal):
+    default: float | None = None
+
+
+class Int64Attribute(BaseModelLocal):
+    default: int | None = None
+
+
+class MapAttribute(BaseModelLocal):
+    default: CustomDefault | None = None
+    element_type: ElemType
+
+
+class NestedAttributeObject(BaseModelLocal):
+    attributes: list[Attribute]
+
+
+class MapNestedAttribute(BaseModelLocal):
+    default: CustomDefault | None = None
+    nested_object: NestedAttributeObject
+
+
+class NumberAttribute(BaseModelLocal):
+    default: CustomDefault | None = None
+
+
+class SetAttribute(BaseModelLocal):
+    default: CustomDefault | None = None
+    element_type: ElemType
+
+
+class SetNestedAttribute(BaseModelLocal):
+    default: CustomDefault | None = None
+    nested_object: NestedAttributeObject
+
+
+class SingleNestedAttribute(BaseModelLocal):
+    default: CustomDefault | None = None
+    nested_object: NestedAttributeObject
+
+
+class StringAttribute(BaseModelLocal):
+    default: str | None = None
+
+
+class ListAttribute(BaseModelLocal):
+    default: CustomDefault | None = None
+    element_type: ElemType
+
+
+class ListNestedAttribute(BaseModelLocal):
+    default: CustomDefault | None = None
+    nested_object: NestedAttributeObject
+
+
+class Operation(int, Enum):
+    CREATE = 0
+    UPDATE = 1
+    READ = 2
+    DELETE = 3
+
+
+class TimeoutsAttribute(BaseModelLocal):
+    configurable_timeouts: list[Operation] = Field(default_factory=list, alias="configurabletimeouts")
+
+
+NestedModelT: TypeAlias = SetNestedAttribute | ListNestedAttribute | MapNestedAttribute | SingleNestedAttribute
+
+
+class Attribute(BaseModelLocal):
+    list: ListAttribute | None = None
+    float64: Float64Attribute | None = None
+    string: StringAttribute | None = None
+    bool: BoolAttribute | None = None
+    map: MapAttribute | None = None
+    number: NumberAttribute | None = None
+    set: SetAttribute | None = None
+    int64: Int64Attribute | None = None
+
+    set_nested: SetNestedAttribute | None = None
+    list_nested: ListNestedAttribute | None = None
+    map_nested: MapNestedAttribute | None = None
+    single_nested: SingleNestedAttribute | None = None
+
+    timeouts: TimeoutsAttribute | None = None
+    description: str | None = None
+    name: SnakeCaseString
+    deprecation_message: str | None = None
+    sensitive: bool = False
+
+    computed_optional_required: ComputedOptionalRequired | None = Field(default=None, alias="computedoptionalrequired")
+
+    @property
+    def name_pascal(self):
+        return pascalize(self.name)
+
+    @property
+    def is_nested(self) -> bool:
+        return any([self.single_nested, self.list_nested, self.map_nested, self.set_nested])
+
+    @property
+    def nested_model(self) -> NestedModelT:
+        assert self.is_nested
+        return self.single_nested or self.list_nested or self.map_nested or self.set_nested  # type: ignore
+
+    @property
+    def is_attribute(self) -> bool:
+        return self.computed_optional_required != ComputedOptionalRequired.unset
+
+    @property
+    def nested_attributes(self) -> list[Attribute]:
+        return self.nested_model.nested_object.attributes
+
+    @property
+    def is_optional(self) -> bool:
+        return self.computed_optional_required in {
+            ComputedOptionalRequired.optional,
+            ComputedOptionalRequired.computed_optional,
+            ComputedOptionalRequired.computed,
+        }
+
+    @property
+    def is_required(self) -> bool:
+        return self.computed_optional_required == ComputedOptionalRequired.required
+
+    @property
+    def go_type(self) -> str:
+        assert not self.is_nested
+        if self.bool:
+            return "bool"
+        if self.float64:
+            return "float64"
+        if self.int64:
+            return "int64"
+        if self.number:
+            return "number"
+        if self.string:
+            return "string"
+        if self.list:
+            raise NotImplementedError(f"list {self.name}")
+        if self.map:
+            if self.map.element_type == ElemType.STRING:
+                return "map[string]string"
+            raise NotImplementedError(f"map {self.name}")
+        if self.set:
+            raise NotImplementedError(f"set {self.name}")
+        raise NotImplementedError(f"unknown type: {self.name}")
+
+    @property
+    def go_type_optional(self) -> str:
+        assert not self.is_nested, "should find go_type_optional for a nested attribute"
+        return f"*{self.go_type}" if self.is_optional else self.go_type
+
+
+class Schema(BaseModelLocal):
+    description: str | None = None
+    deprecation_message: str | None = None
+    attributes: list[Attribute]
+
+
+class Resource(BaseModelLocal):
+    schema: Schema
+    name: SnakeCaseString
+
+    @property
+    def use_timeout(self) -> bool:
+        return any(a.timeouts for a in self.schema.attributes)
+
+
+ResourceSchemaV3 = Resource
+TF_MODEL_NAME = "TFModel"

atlas_init/cli_tf/schema_v3_sdk.py

@@ -0,0 +1,279 @@
+import logging
+from functools import singledispatch
+from queue import Queue
+
+from atlas_init.cli_tf.schema_v2 import (
+    extend_import_urls,
+    go_fmt,
+    import_lines,
+    package_name,
+)
+from atlas_init.cli_tf.schema_v2_sdk import GoVarName, SDKAttribute, SDKModel
+from atlas_init.cli_tf.schema_v3 import (
+    TF_MODEL_NAME,
+    Attribute,
+    ListNestedAttribute,
+    Resource,
+    SingleNestedAttribute,
+)
+from atlas_init.cli_tf.schema_v3_sdk_base import (
+    AllowedMissingAttributeError,
+    SDKAndSchemaAttribute,
+    find_attribute,
+    name_schema_struct,
+    name_struct_attribute,
+    schema_attributes,
+)
+from atlas_init.humps import camelize, pascalize
+
+logger = logging.getLogger(__name__)
+
+
+def generate_model_go(resource: Resource, sdk_model: SDKModel) -> str:
+    func_lines = sdk_to_tf_func(resource, sdk_model)
+    import_urls = set()
+    extend_import_urls(import_urls, func_lines)
+    unformatted = "\n".join(
+        [
+            f"package {package_name(resource.name)}",
+            "",
+            *import_lines(import_urls),
+            "",
+            *func_lines,
+        ]
+    )
+    return go_fmt(resource.name, unformatted)
+
+
+def sdk_to_tf_func(resource: Resource, sdk_model: SDKModel) -> list[str]:
+    lines = []
+    timeouts_signature = ", timeout timeouts.Value" if resource.use_timeout else ""
+    lines.append(
+        f"func New{TF_MODEL_NAME}(ctx context.Context, {GoVarName.INPUT} *admin.{sdk_model.name}{timeouts_signature}) (*{TF_MODEL_NAME}, diag.Diagnostics) {{"
+    )
+    nested_attributes, call_lines = call_nested_functions(resource, sdk_model.list_nested_attributes())
+    lines.extend(call_lines)
+    timeouts_set = ["    Timeouts: timeout,"] if resource.use_timeout else []
+    lines.extend(
+        [
+            f"  return &{TF_MODEL_NAME}{{",
+            *tf_struct_create(resource, sdk_model),
+            *timeouts_set,
+            "  }, nil",  # close return
+            "}\n",  # close function
+        ]
+    )
+    lines.extend(process_nested_attributes(nested_attributes))
+    return lines
+
+
+def as_var_name(attr: SDKAttribute) -> str:
+    return camelize(attr.json_name)
+
+
+def call_nested_functions(
+    root: Resource | Attribute,
+    nested_attributes: list[SDKAttribute],
+    *,
+    return_on_error: bool = True,
+    sdk_var_name: GoVarName = GoVarName.INPUT,
+) -> tuple[list[SDKAndSchemaAttribute], list[str]]:
+    lines = []
+    schema_nested_attributes = schema_attributes(root)
+    nested_generations = []
+    for sdk_attribute in nested_attributes:
+        try:
+            schema_attribute = find_attribute(schema_nested_attributes, sdk_attribute.tf_name, root.name)
+        except AllowedMissingAttributeError as e:
+            logger.info(f"skipping {e!r}")
+            continue
+        var_name = as_var_name(sdk_attribute)
+        lines.append(
+            f"{var_name} := New{_name_custom_object_type(schema_attribute.name)}(ctx, {sdk_var_name}.{sdk_attribute.struct_name}, {GoVarName.DIAGS})"
+        )
+        nested_generations.append(SDKAndSchemaAttribute(sdk_attribute, schema_attribute))
+    if lines and return_on_error:
+        lines.insert(0, "diags := &diag.Diagnostics{}")
+        lines.extend(
+            [
+                "  if diags.HasError() {",
+                "    return nil, *diags",
+                "  }",
+            ]
+        )
+
+    return nested_generations, lines
+
+
+_sdk_to_tf_funcs = {
+    ("*string", "string"): lambda sdk_ref: f"types.StringPointerValue({sdk_ref})",
+    ("string", "string"): lambda sdk_ref: f"types.StringValue({sdk_ref})",
+    ("*int64", "int64"): lambda sdk_ref: f"types.Int64PointerValue({sdk_ref})",
+    ("int64", "int64"): lambda sdk_ref: f"types.Int64Value({sdk_ref})",
+    (
+        "*int",
+        "int64",
+    ): lambda sdk_ref: f"types.Int64PointerValue(conversion.IntPtrToInt64Ptr({sdk_ref}))",
+    ("*float64", "float64"): lambda sdk_ref: f"types.Float64PointerValue({sdk_ref})",
+    ("float64", "float64"): lambda sdk_ref: f"types.Float64Value({sdk_ref})",
+    ("*bool", "bool"): lambda sdk_ref: f"types.BoolPointerValue({sdk_ref})",
+    ("bool", "bool"): lambda sdk_ref: f"types.BoolValue({sdk_ref})",
+    (
+        "*map[string]string",
+        "map[string]string",
+    ): lambda sdk_ref: f"conversion.ToTFMapOfString({GoVarName.CTX}, {GoVarName.DIAGS}, {sdk_ref})",
+    (
+        "map[string]string",
+        "map[string]string",
+    ): lambda sdk_ref: f"conversion.ToTFMapOfString({GoVarName.CTX}, {GoVarName.DIAGS}, &{sdk_ref})",
+    (
+        "*time.Time",
+        "string",
+    ): lambda sdk_ref: f"types.StringPointerValue(conversion.TimePtrToStringPtr({sdk_ref}))",
+    (
+        "time.Time",
+        "string",
+    ): lambda sdk_ref: f"types.StringValue(conversion.TimeToString({sdk_ref}))",
+}
+
+
+def sdk_to_tf_attribute_value(
+    schema_attribute: Attribute,
+    sdk_attribute: SDKAttribute,
+    variable_name: GoVarName = GoVarName.INPUT,
+) -> str:
+    key = (sdk_attribute.go_type, schema_attribute.go_type)
+    if key in _sdk_to_tf_funcs:
+        return _sdk_to_tf_funcs[key](f"{variable_name}.{sdk_attribute.struct_name}")
+    raise ValueError(f"Could not find conversion function for {key}")
+
+
+# sdk_to_tf_attribute_value(schema_attribute, sdk_attribute, sdk_var_name)
+def tf_struct_create(
+    root: Resource | Attribute,
+    sdk_model: SDKModel,
+    sdk_var_name: GoVarName = GoVarName.INPUT,
+) -> list[str]:
+    lines = []
+    for attr in schema_attributes(root):
+        if attr.is_nested:
+            local_var = sdk_model.lookup_tf_name(attr.name)
+            lines.append(f"{name_struct_attribute(attr.name)}: {as_var_name(local_var)},")
+        elif attr.is_attribute:
+            local_var = sdk_model.lookup_tf_name(attr.name)
+            lines.append(
+                f"{name_struct_attribute(attr.name)}: {sdk_to_tf_attribute_value(attr, local_var, sdk_var_name)},"
+            )
+    return lines
+
+
+def process_nested_attributes(
+    nested_attributes: list[SDKAndSchemaAttribute],
+) -> list[str]:
+    lines = []
+    queue = Queue()
+
+    def add_nested_to_queue(attributes: list[SDKAndSchemaAttribute]):
+        for nested in attributes:
+            logger.info(f"found nested attribute: {nested.schema_attribute.name}")
+            queue.put(nested)
+
+    add_nested_to_queue(nested_attributes)
+    while not queue.empty():
+        sdk_attribute, schema_attribute = queue.get()
+        more_nested_attributes, nested_lines = convert_nested_attribute(
+            schema_attribute.nested_model, schema_attribute, sdk_attribute
+        )
+        lines.extend(nested_lines)
+        add_nested_to_queue(more_nested_attributes)
+    return lines
+
+
+def _name_custom_object_type(name: str) -> str:
+    return f"{pascalize(name)}ObjType"
+
+
+@singledispatch
+def convert_nested_attribute(
+    nested_model: object, schema_attribute: Attribute, _: SDKAttribute
+) -> tuple[list[SDKAndSchemaAttribute], list[str]]:
+    raise NotImplementedError(f"unsupported nested attribute: {schema_attribute.name} of type {type(nested_model)}")
+
+
+@convert_nested_attribute.register
+def _convert_single_nested_attribute(
+    _: SingleNestedAttribute,
+    schema_attribute: Attribute,
+    sdk_attribute: SDKAttribute,
+) -> tuple[list[SDKAndSchemaAttribute], list[str]]:
+    object_type_name = _name_custom_object_type(schema_attribute.name)
+    lines: list[str] = [
+        f"func New{object_type_name}(ctx context.Context, {GoVarName.INPUT} *admin.{sdk_attribute.struct_type_name}, diags *diag.Diagnostics) types.Object {{",
+        f"  if {GoVarName.INPUT} == nil {{",
+        f"    return types.ObjectNull({object_type_name}.AttrTypes)",
+        "  }",
+    ]
+
+    nested_attributes, call_lines = call_nested_functions(
+        schema_attribute, sdk_attribute.list_nested_attributes(), return_on_error=False
+    )
+    lines.extend(call_lines)
+    struct_name = name_schema_struct(schema_attribute.name)
+    lines.extend(
+        [
+            f"  tfModel := {struct_name}{{",
+            *tf_struct_create(schema_attribute, sdk_attribute.as_sdk_model()),
+            "  }",
+            f"  objType, diagsLocal := types.ObjectValueFrom(ctx, {object_type_name}.AttrTypes, tfModel)",
+            f"  {GoVarName.DIAGS}.Append(diagsLocal...)",
+            "  return objType",
+            "}\n",
+        ]
+    )
+    return nested_attributes, lines
+
+
+@convert_nested_attribute.register
+def _convert_list_nested_attriute(
+    _: ListNestedAttribute,
+    schema_attribute: Attribute,
+    sdk_attribute: SDKAttribute,
+) -> tuple[list[SDKAndSchemaAttribute], list[str]]:
+    object_type_name = _name_custom_object_type(schema_attribute.name)
+    lines: list[str] = [
+        f"func New{object_type_name}(ctx context.Context, {GoVarName.INPUT} *[]admin.{sdk_attribute.struct_type_name}, diags *diag.Diagnostics) types.List {{",
+        f"  if {GoVarName.INPUT} == nil {{",
+        f"    return types.ListNull({object_type_name})",
+        "  }",
+    ]
+    struct_name = name_schema_struct(schema_attribute.name)
+    lines.extend(
+        [
+            f"  tfModels := make([]{struct_name}, len(*{GoVarName.INPUT}))",
+            f"  for i, item := range *{GoVarName.INPUT} {{",
+        ]
+    )
+    nested_attributes, call_lines = call_nested_functions(
+        schema_attribute,
+        sdk_attribute.list_nested_attributes(),
+        return_on_error=False,
+        sdk_var_name=GoVarName.ITEM,
+    )
+    lines.extend([f"  {line}" for line in call_lines])
+    lines.extend(
+        [
+            f"    tfModels[i] = {struct_name}{{",
+            *tf_struct_create(
+                schema_attribute,
+                sdk_attribute.as_sdk_model(),
+                sdk_var_name=GoVarName.ITEM,
+            ),
+            "    }",
+            "  }",
+            f"  listType, diagsLocal := types.ListValueFrom(ctx, {object_type_name}, tfModels)",
+            f"  {GoVarName.DIAGS}.Append(diagsLocal...)",
+            "  return listType",
+            "}\n",
+        ]
+    )
+    return nested_attributes, lines

atlas_init/cli_tf/schema_v3_sdk_base.py

@@ -0,0 +1,68 @@
+from functools import singledispatch
+from typing import NamedTuple
+
+from atlas_init.cli_tf.schema_v2_sdk import SDKAttribute
+from atlas_init.cli_tf.schema_v3 import Attribute, Resource
+from atlas_init.humps import pascalize
+
+
+class SDKAndSchemaAttribute(NamedTuple):
+    sdk_attribute: SDKAttribute
+    schema_attribute: Attribute
+
+
+_allowed_missing_attributes: set[str] = set()
+
+
+def set_allowed_missing_attributes(allowed: set[str]):
+    global _allowed_missing_attributes  # noqa: PLW0603 `Using the global statement to update `_allowed_missing_attributes` is discouraged`
+    _allowed_missing_attributes = allowed
+
+
+class AllowedMissingAttributeError(Exception):
+    def __init__(self, name: str, *args: object) -> None:
+        self.name = name
+        super().__init__(*args)
+
+
+def find_attribute(attributes: list[Attribute], name: str, root_name: str) -> Attribute:
+    for schema_attribute in attributes:
+        if name == schema_attribute.name:
+            return schema_attribute
+    if name in _allowed_missing_attributes:
+        raise AllowedMissingAttributeError(name)
+    raise ValueError(f"could not find schema attribute for {name} on resource: {root_name}")
+
+
+@singledispatch
+def schema_attributes(root: object) -> list[Attribute]:
+    raise NotImplementedError(f"unsupported root type: {type(root)} to find schema attributes")
+
+
+@schema_attributes.register
+def _resource_attributes(root: Resource) -> list[Attribute]:
+    return root.schema.attributes
+
+
+@schema_attributes.register
+def _attribute_nested(root: Attribute) -> list[Attribute]:
+    return root.nested_attributes
+
+
+def name_schema_struct(name: str) -> str:
+    return f"TF{pascalize(name)}Model"
+
+
+def name_struct_attribute(name: str) -> str:
+    default = pascalize(name)
+    if override := _name_attribute_overrides.get(default):
+        return override
+    return default
+
+
+_name_attribute_overrides = {}
+
+
+def set_name_attribute_overrides(overrides: dict[str, str]):
+    global _name_attribute_overrides  # noqa: PLW0603 `Using the global statement to update `_name_attribute_overrides` is discouraged`
+    _name_attribute_overrides = overrides