atlas-init 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

atlas_init/tf_ext/settings.py

@@ -0,0 +1,188 @@
+import logging
+from pathlib import Path
+import re
+from typing import ClassVar, Self
+
+from model_lib import Entity, StaticSettings
+from pydantic import model_validator
+from zero_3rdparty.file_utils import ensure_parents_write_text
+from zero_3rdparty.str_utils import ensure_suffix
+
+from atlas_init.tf_ext.args import ENV_NAME_REPO_PATH_ATLAS_PROVIDER, TF_CLI_CONFIG_FILE_ENV_NAME
+
+logger = logging.getLogger(__name__)
+
+_atlas_provider_installation = re.compile(r'"mongodb/mongodbatlas"\s*=\s*"(?P<repo_path>[^"]+)/bin"')
+
+
+def parse_atlas_path_from_tf_cli_config_file(tf_cli_config_file: Path | None) -> Path | None:
+    assert tf_cli_config_file, "tf_cli_config_file is required to parse atlas path"
+    text = tf_cli_config_file.read_text()
+    if match := _atlas_provider_installation.search(text):
+        return Path(match.group("repo_path"))
+    return None
+
+
+_tfrc_template = """\
+provider_installation {
+  dev_overrides {
+    "mongodb/mongodbatlas" = "REPO_PATH_TF_PROVIDER"
+  }
+  direct {}
+}
+"""
+
+
+def tf_cli_config_file_content(repo_path_atlas: Path | None) -> str:
+    assert repo_path_atlas, "repo_path_atlas is required to export tf_cli_config_file"
+    return _tfrc_template.replace("REPO_PATH_TF_PROVIDER", ensure_suffix(str(repo_path_atlas), "/bin"))
+
+
+def resource_type_name_no_provider(provider_name: str, resource_type: str) -> str:
+    return resource_type.removeprefix(provider_name).strip("_")
+
+
+class RepoOut(Entity):
+    base: Path
+
+    @property
+    def resource_modules(self) -> Path:
+        return self.base / "resource_modules"
+
+    def resource_modules_provider_path(self, provider_name: str) -> Path:
+        return self.resource_modules / provider_name
+
+    def resource_module_path(self, provider_name: str, resource_type: str) -> Path:
+        return self.resource_modules_provider_path(provider_name) / resource_type_name_no_provider(
+            provider_name, resource_type
+        )
+
+    def dataclass_path(self, provider_name: str, resource_type: str) -> Path:
+        return self.py_provider_module(provider_name) / f"{resource_type}.py"
+
+    @property
+    def py_modules(self) -> Path:
+        return self.base / "py_modules"
+
+    @property
+    def config_path(self) -> Path:
+        return self.base / "config"
+
+    def py_provider_module(self, provider: str) -> Path:
+        return self.py_modules / f"tf_{provider}"
+
+    def provider_settings_path(self, provider: str) -> Path:
+        return self.config_path / f"{provider}.yaml"
+
+
+class TfExtSettings(StaticSettings):
+    ENV_NAME_REPO_PATH_ATLAS_PROVIDER: ClassVar[str] = ENV_NAME_REPO_PATH_ATLAS_PROVIDER
+    ENV_NAME_TF_CLI_CONFIG_FILE: ClassVar[str] = TF_CLI_CONFIG_FILE_ENV_NAME
+
+    repo_path_atlas_provider: Path | None = None
+    tf_cli_config_file: Path | None = None
+    repo_out_path: Path | None = None
+
+    @model_validator(mode="after")
+    def infer_repo_path_atlas(self) -> Self:
+        if self.repo_path_atlas_provider is None and self.tf_cli_config_file is None:
+            raise ValueError("repo_path_atlas or tf_cli_config_file must be set")
+        if self.repo_path_atlas_provider is None:
+            self.repo_path_atlas_provider = parse_atlas_path_from_tf_cli_config_file(self.tf_cli_config_file)
+        if self.tf_cli_config_file is None:
+            cli_config_file_content = tf_cli_config_file_content(self.repo_path_atlas_provider)
+            self.tf_cli_config_file = self.static_root / "dev.tfrc"
+            ensure_parents_write_text(self.tf_cli_config_file, cli_config_file_content)
+        if self.tf_cli_config_file:
+            tf_cli_repo_path = parse_atlas_path_from_tf_cli_config_file(self.tf_cli_config_file)
+            assert tf_cli_repo_path == self.repo_path_atlas_provider, (
+                f"tf_cli_config_file does not match repo_path_atlas_provider {tf_cli_repo_path} != {self.repo_path_atlas_provider}"
+            )
+        return self
+
+    @property
+    def repo_out(self) -> RepoOut:
+        assert self.repo_out_path, "repo_out_path is required"
+        return RepoOut(base=self.repo_out_path)
+
+    @property
+    def atlas_graph_path(self) -> Path:
+        return self.static_root / "atlas_graph.yaml"
+
+    @property
+    def vars_file_path(self) -> Path:
+        return self.static_root / "tf_vars.yaml"
+
+    @property
+    def vars_external_file_path(self) -> Path:
+        return self.static_root / "tf_vars_external.yaml"
+
+    @property
+    def resource_types_file_path(self) -> Path:
+        return self.static_root / "tf_resource_types.yaml"
+
+    @property
+    def resource_types_external_file_path(self) -> Path:
+        return self.static_root / "tf_resource_types_external.yaml"
+
+    @property
+    def schema_resource_types_path(self) -> Path:
+        return self.static_root / "tf_schema_resource_types.yaml"
+
+    @property
+    def schema_resource_types_deprecated_path(self) -> Path:
+        return self.static_root / "tf_schema_resource_types_deprecated.yaml"
+
+    @property
+    def api_calls_path(self) -> Path:
+        return self.static_root / "tf_api_calls.yaml"
+
+    def pagination_output_path(self, query_string: str) -> Path:
+        return self.static_root / "pagination_output" / f"query_is_{query_string or 'empty'}.md"
+
+    @property
+    def new_res_path(self) -> Path:
+        return self.static_root / "newres"
+
+    @property
+    def modules_out_path(self) -> Path:
+        return self.static_root / "modules"
+
+    @property
+    def attribute_description_file_path(self) -> Path:
+        return self.static_root / "attribute_description.yaml"
+
+    @property
+    def attribute_description_manual_file_path(self) -> Path:
+        return self.static_root / "attribute_description_manual.yaml"
+
+    @property
+    def attribute_resource_descriptions_file_path(self) -> Path:
+        return self.static_root / "attribute_resource_descriptions.yaml"
+
+    @property
+    def attribute_resource_descriptions_manual_file_path(self) -> Path:
+        return self.static_root / "attribute_resource_descriptions_manual.yaml"
+
+    @property
+    def output_plan_dumps(self) -> Path:
+        return self.static_root / "output_plan_dumps"
+
+    @property
+    def plan_diff_output_path(self) -> Path:
+        return self.static_root / "plan_diff_output"
+
+    def provider_cache_dir(self, provider_name: str) -> Path:
+        return self.cache_root / "provider_cache" / provider_name
+
+
+def init_tf_ext_settings(*, allow_empty_out_path: bool = False) -> TfExtSettings:
+    settings = TfExtSettings.from_env()
+    assert settings
+    logger.info("env-vars ready: ✅")
+    logger.info(f"repo_path_atlas: {settings.repo_path_atlas_provider}")
+    logger.info(f"tf_cli_config_file: {settings.tf_cli_config_file}")
+    if not allow_empty_out_path:
+        assert settings.repo_out
+    logger.info(f"Repo out path is: {settings.repo_out_path}")
+    return settings

atlas_init/tf_ext/tf_dep.py

@@ -0,0 +1,324 @@
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from pathlib import Path
+from typing import Iterable, NamedTuple
+
+import pydot
+from ask_shell import ShellError, new_task, run_and_wait
+from ask_shell._run import stop_runs_and_pool
+from ask_shell.run_pool import run_pool
+from model_lib import Entity, dump
+from pydantic import BaseModel, Field
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed
+from typer import Typer
+from zero_3rdparty.file_utils import ensure_parents_write_text
+from zero_3rdparty.iter_utils import flat_map
+
+from atlas_init.settings.rich_utils import configure_logging
+from atlas_init.tf_ext.args import REPO_PATH_ATLAS_ARG, SKIP_EXAMPLES_DIRS_OPTION
+from atlas_init.tf_ext.constants import ATLAS_PROVIDER_NAME
+from atlas_init.tf_ext.paths import find_variable_resource_type_usages, find_variables, get_example_directories
+from atlas_init.tf_ext.settings import TfExtSettings
+
+logger = logging.getLogger(__name__)
+v2_grand_parent_dirs = {
+    "module_maintainer",
+    "module_user",
+    "migrate_cluster_to_advanced_cluster",
+    "mongodbatlas_backup_compliance_policy",
+}
+v2_parent_dir = {"cluster_with_schedule"}
+MODULE_PREFIX = "module."
+DATA_PREFIX = "data."
+VARIABLE_RESOURCE_MAPPING: dict[str, str] = {
+    "org_id": "mongodbatlas_organization",
+    "project_id": "mongodbatlas_project",
+    "cluster_name": "mongodbatlas_advanced_cluster",
+}
+SKIP_NODES: set[str] = {"mongodbatlas_cluster", "mongodbatlas_flex_cluster"}
+FORCE_INTERNAL_NODES: set[str] = {"mongodbatlas_project_ip_access_list"}
+
+
+def is_v2_example_dir(example_dir: Path) -> bool:
+    parent_dir = example_dir.parent.name
+    grand_parent_dir = example_dir.parent.parent.name
+    return parent_dir in v2_parent_dir or grand_parent_dir in v2_grand_parent_dirs
+
+
+def tf_dep_graph(
+    repo_path: Path = REPO_PATH_ATLAS_ARG,
+    skip_names: list[str] = SKIP_EXAMPLES_DIRS_OPTION,
+):
+    settings = TfExtSettings.from_env()
+    output_dir = settings.static_root
+    logger.info(f"Using output directory: {output_dir}")
+    example_dirs = get_example_directories(repo_path, skip_names)
+    logger.info(f"example_dirs: \n{'\n'.join(str(d) for d in sorted(example_dirs))}")
+    with new_task("Find terraform graphs", total=len(example_dirs)) as task:
+        atlas_graph = parse_graphs(example_dirs, task)
+    with new_task("Dump graph"):
+        graph_yaml = atlas_graph.dump_yaml()
+        ensure_parents_write_text(settings.atlas_graph_path, graph_yaml)
+        logger.info(f"Atlas graph dumped to {settings.atlas_graph_path}")
+
+
+def print_edges(graph: pydot.Dot):
+    edges = graph.get_edges()
+    for edge in edges:
+        logger.info(f"{edge.get_source()} -> {edge.get_destination()}")
+
+
+class ResourceParts(NamedTuple):
+    resource_type: str
+    resource_name: str
+
+    @property
+    def provider_name(self) -> str:
+        return self.resource_type.split("_")[0]
+
+
+class ResourceRef(BaseModel):
+    full_ref: str
+
+    def _resource_parts(self) -> ResourceParts:
+        match self.full_ref.split("."):
+            case [resource_type, resource_name] if "_" in resource_type:
+                return ResourceParts(resource_type, resource_name)
+            case [*_, resource_type, resource_name] if "_" in resource_type:
+                return ResourceParts(resource_type, resource_name)
+        raise ValueError(f"Invalid resource reference: {self.full_ref}")
+
+    @property
+    def provider_name(self) -> str:
+        return self._resource_parts().provider_name
+
+    @property
+    def is_external(self) -> bool:
+        return self.provider_name != ATLAS_PROVIDER_NAME
+
+    @property
+    def is_atlas_resource(self) -> bool:
+        return not self.is_module and not self.is_data and self.provider_name == ATLAS_PROVIDER_NAME
+
+    @property
+    def is_module(self) -> bool:
+        return self.full_ref.startswith(MODULE_PREFIX)
+
+    @property
+    def is_data(self) -> bool:
+        return self.full_ref.startswith(DATA_PREFIX)
+
+    @property
+    def resource_type(self) -> str:
+        return self._resource_parts().resource_type
+
+
+class EdgeParsed(BaseModel):
+    parent: ResourceRef
+    child: ResourceRef
+
+    @classmethod
+    def from_edge(cls, edge: pydot.Edge) -> "EdgeParsed":
+        return cls(
+            # edges shows from child --> parent, so we reverse the order
+            parent=ResourceRef(full_ref=edge_plain(edge.get_destination())),
+            child=ResourceRef(full_ref=edge_plain(edge.get_source())),
+        )
+
+    @property
+    def has_module_edge(self) -> bool:
+        return self.parent.is_module or self.child.is_module
+
+    @property
+    def has_data_edge(self) -> bool:
+        return self.parent.is_data or self.child.is_data
+
+    @property
+    def is_resource_edge(self) -> bool:
+        return not self.has_module_edge and not self.has_data_edge
+
+    @property
+    def is_external_to_internal_edge(self) -> bool:
+        return self.parent.is_external and self.child.is_atlas_resource
+
+    @property
+    def is_internal_atlas_edge(self) -> bool:
+        return self.parent.is_atlas_resource and self.child.is_atlas_resource
+
+
+def edge_plain(edge_endpoint: pydot.EdgeEndpoint) -> str:
+    return str(edge_endpoint).strip('"').strip()
+
+
+def edge_src_dest(edge: pydot.Edge) -> tuple[str, str]:
+    """Get the source and destination of the edge as plain strings."""
+    return edge_plain(edge.get_source()), edge_plain(edge.get_destination())
+
+
+def skip_variable_edge(src: str, dst: str) -> bool:
+    # sourcery skip: assign-if-exp, boolean-if-exp-identity, reintroduce-else, remove-unnecessary-cast
+    if src == dst:
+        return True
+    if src == "mongodbatlas_advanced_cluster" and "_cluster" in dst:
+        return True
+    return False
+
+
+class AtlasGraph(Entity):
+    # atlas_resource_type -> set[atlas_resource_type]
+    parent_child_edges: dict[str, set[str]] = Field(default_factory=lambda: defaultdict(set))
+    # atlas_resource_type -> set[external_resource_type]
+    external_parents: dict[str, set[str]] = Field(default_factory=lambda: defaultdict(set))
+    deprecated_resource_types: set[str] = Field(default_factory=set)
+
+    def all_parents(self, child: str) -> Iterable[str]:
+        for parent, children in self.parent_child_edges.items():
+            if child in children:
+                yield parent
+
+    def dump_yaml(self) -> str:
+        parent_child_edges = {name: sorted(children) for name, children in sorted(self.parent_child_edges.items())}
+        external_parents = {name: sorted(parents) for name, parents in sorted(self.external_parents.items())}
+        return dump(
+            {
+                "parent_child_edges": parent_child_edges,
+                "external_parents": external_parents,
+            },
+            format="yaml",
+        )
+
+    @property
+    def all_internal_nodes(self) -> set[str]:
+        return set(flat_map([src] + list(dsts) for src, dsts in self.parent_child_edges.items()))
+
+    def iterate_internal_edges(self) -> Iterable[tuple[str, str]]:
+        for parent, children in self.parent_child_edges.items():
+            for child in children:
+                yield parent, child
+
+    @property
+    def all_external_nodes(self) -> set[str]:
+        return set(flat_map([src] + list(dsts) for src, dsts in self.external_parents.items()))
+
+    def iterate_external_edges(self) -> Iterable[tuple[str, str]]:
+        for child, parents in self.external_parents.items():
+            for parent in parents:
+                yield parent, child
+
+    def add_edges(self, edges: list[pydot.Edge]):
+        for edge in edges:
+            parsed = EdgeParsed.from_edge(edge)
+            parent = parsed.parent
+            child = parsed.child
+            if parsed.is_internal_atlas_edge:
+                self.parent_child_edges[parent.resource_type].add(child.resource_type)
+                # edges shows from child --> parent, so we reverse the order
+            elif parsed.is_external_to_internal_edge:
+                if parent.provider_name in {"random", "cedar"}:
+                    continue  # skip random provider edges
+                self.external_parents[child.resource_type].add(parent.resource_type)
+
+    def add_variable_edges(self, example_dir: Path) -> None:
+        """Use the variables to find the resource dependencies."""
+        if not (variables := find_variables(example_dir / "variables.tf")):
+            return
+        usages = find_variable_resource_type_usages(set(variables), example_dir)
+        for variable, resource_types in usages.items():
+            if parent_type := VARIABLE_RESOURCE_MAPPING.get(variable):
+                for child_type in resource_types:
+                    if skip_variable_edge(parent_type, child_type):
+                        continue
+                    if child_type.startswith(ATLAS_PROVIDER_NAME):
+                        logger.info(f"Adding variable edge: {parent_type} -> {child_type}")
+                        self.parent_child_edges[parent_type].add(child_type)
+
+
+def parse_graphs(example_dirs: list[Path], task: new_task, max_workers: int = 16, max_dirs: int = 9999) -> AtlasGraph:
+    atlas_graph = AtlasGraph()
+    with run_pool("parse example graphs", total=len(example_dirs)) as executor:
+        futures = {
+            executor.submit(parse_graph, example_dir): example_dir
+            for i, example_dir in enumerate(example_dirs)
+            if i < max_dirs
+        }
+        graphs = {}
+        for future in futures:
+            try:
+                example_dir, graph_output = future.result()
+            except ShellError as e:
+                logger.error(f"Error parsing graph for {futures[future]}: {e}")
+                continue
+            except KeyboardInterrupt:
+                logger.error("KeyboardInterrupt received, stopping graph parsing.")
+                stop_runs_and_pool("KeyboardInterrupt", immediate=True)
+                break
+            try:
+                graph = graphs[example_dir] = parse_graph_output(example_dir, graph_output)
+            except GraphParseError as e:
+                logger.error(e)
+                continue
+            atlas_graph.add_edges(graph.get_edges())
+            atlas_graph.add_variable_edges(example_dir)
+            task.update(advance=1)
+    return atlas_graph
+
+
+class GraphParseError(Exception):
+    def __init__(self, example_dir: Path, message: str):
+        self.example_dir = example_dir
+        super().__init__(f"Failed to parse graph for {example_dir}: {message}")
+
+
+def parse_graph_output(example_dir: Path, graph_output: str, verbose: bool = False) -> pydot.Dot:
+    assert graph_output, f"Graph output is empty for {example_dir}"
+    dots = pydot.graph_from_dot_data(graph_output)  # not thread safe, so we use the main thread here instead
+    if not dots:
+        raise GraphParseError(example_dir, f"No graphs found in the output:\n{graph_output}")
+    assert len(dots) == 1, f"Expected one graph for {example_dir}, got {len(dots)}"
+    graph = dots[0]
+    edges = graph.get_edges()
+    if not edges:
+        logger.info(f"No edges found in graph for {example_dir}")
+    if verbose:
+        print_edges(graph)
+    return graph
+
+
+class EmptyGraphOutputError(Exception):
+    """Raised when the graph output is empty."""
+
+    def __init__(self, example_dir: Path):
+        self.example_dir = example_dir
+        super().__init__(f"Graph output is empty for {example_dir}")
+
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_fixed(1),
+    retry=retry_if_exception_type(EmptyGraphOutputError),
+    reraise=True,
+)
+def parse_graph(example_dir: Path) -> tuple[Path, str]:
+    env_vars = {
+        "MONGODB_ATLAS_PREVIEW_PROVIDER_V2_ADVANCED_CLUSTER": "true" if is_v2_example_dir(example_dir) else "false",
+    }
+    lock_file = example_dir / ".terraform.lock.hcl"
+    if not lock_file.exists():
+        run_and_wait("terraform init", cwd=example_dir, env=env_vars)
+    run = run_and_wait("terraform graph", cwd=example_dir, env=env_vars)
+    if graph_output := run.stdout_one_line:
+        return example_dir, graph_output
+    raise EmptyGraphOutputError(example_dir)
+
+
+def typer_main():
+    app = Typer()
+    app.command()(tf_dep_graph)
+    configure_logging(app)
+    app()
+
+
+if __name__ == "__main__":
+    typer_main()

atlas_init/tf_ext/tf_desc_gen.py

@@ -0,0 +1,53 @@
+import logging
+from collections import defaultdict
+from atlas_init.tf_ext.args import TF_CLI_CONFIG_FILE_ARG
+from atlas_init.tf_ext.settings import TfExtSettings
+from atlas_init.tf_ext.provider_schema import ResourceSchema, parse_atlas_schema
+from model_lib import dump, parse_model
+from zero_3rdparty.file_utils import ensure_parents_write_text
+
+logger = logging.getLogger(__name__)
+
+
+def tf_desc_gen(
+    tf_cli_config_file: str = TF_CLI_CONFIG_FILE_ARG,
+):
+    settings = TfExtSettings.from_env()
+    out_path = settings.attribute_description_file_path
+    resource_out_path = settings.attribute_resource_descriptions_file_path
+    assert tf_cli_config_file
+    schema = parse_atlas_schema()
+    descriptions = {}
+    descriptions_by_resource: dict[str, dict[str, str]] = defaultdict(dict)
+    attr_desc_resource = {}
+
+    def add_description(resource_type: str, attr_name: str, description: str | None) -> None:
+        if not description:
+            return
+        descriptions_by_resource[resource_type][attr_name] = description
+        if existing := descriptions.get(attr_name):
+            if existing != description:
+                old_resource_type = attr_desc_resource[attr_name]
+                logger.info(
+                    f"Descriptions differs between '{old_resource_type}' and '{resource_type}' for attribute '{attr_name}':\n{existing}\n{description}"
+                )
+                if len(existing) > len(description):
+                    return
+        descriptions[attr_name] = description
+        attr_desc_resource[attr_name] = resource_type
+
+    for resource_type, resource_schema in schema.raw_resource_schema.items():
+        parsed_schema = parse_model(resource_schema, t=ResourceSchema)
+        schema_block = parsed_schema.block
+        for name, attribute in (schema_block.attributes or {}).items():
+            add_description(resource_type, name, attribute.description)
+        for name, block_type in (schema_block.block_types or {}).items():
+            add_description(resource_type, name, block_type.description)
+    descriptions_yaml = dump(dict(sorted(descriptions.items())), format="yaml")
+    ensure_parents_write_text(out_path, descriptions_yaml)
+    logger.info(f"Generated attribute descriptions to {out_path}")
+    resource_descriptions_yaml = dump(
+        {k: dict(sorted(v.items())) for k, v in sorted(descriptions_by_resource.items())}, format="yaml"
+    )
+    ensure_parents_write_text(resource_out_path, resource_descriptions_yaml)
+    logger.info(f"Generated attribute resource descriptions to {resource_out_path}")