From a69edcc2c857f7ef3d83c2f3ca7536e881c4cd0b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?= Date: Thu, 2 Sep 2021 14:09:31 +0000 Subject: [PATCH] tpi: lazily import python_terraform (#6526) * machine: lazy load tpi/python_terraform only when needed * remove terraform from package installers * update tpi test imports --- .github/workflows/packages.yaml | 10 +++--- dvc/machine/backend/terraform.py | 52 ++++------------------------ dvc/tpi/__init__.py | 47 +++++++++++++++++++++++++ tests/unit/machine/test_terraform.py | 2 +- 4 files changed, 60 insertions(+), 51 deletions(-) diff --git a/.github/workflows/packages.yaml b/.github/workflows/packages.yaml index 6d225d9fe..f97e85cee 100644 --- a/.github/workflows/packages.yaml +++ b/.github/workflows/packages.yaml @@ -20,7 +20,7 @@ jobs: ruby-version: 2.4 - name: Install deps run: | - pip install .[all,terraform] + pip install .[all] pip install -r scripts/build-requirements.txt gem install --no-document fpm - name: Build deb @@ -49,7 +49,7 @@ jobs: ruby-version: 2.4 - name: Install deps run: | - pip install .[all,terraform] + pip install .[all] pip install -r scripts/build-requirements.txt gem install --no-document fpm - name: Build rpm @@ -81,7 +81,7 @@ jobs: # available lxml (webdav dependency) wheels are built with an old # SDK, which breaks osxpkg notarization. Building from scratch # instead. - pip install .[all,terraform] --no-binary lxml + pip install .[all] --no-binary lxml pip install -r scripts/build-requirements.txt gem install --no-document fpm - name: Build @@ -106,7 +106,7 @@ jobs: python-version: 3.8 - name: Install deps run: | - pip install .[all,terraform] + pip install .[all] pip install -r scripts/build-requirements.txt - name: Pull images run: | @@ -133,7 +133,7 @@ jobs: with: python-version: 3.7 - name: Install - run: pip install ".[all,terraform,tests]" + run: pip install ".[all,tests]" - name: Build packages run: ./scripts/build_package.sh - name: Publish packages to PyPI diff --git a/dvc/machine/backend/terraform.py b/dvc/machine/backend/terraform.py index 7accbb705..43b89ddf9 100644 --- a/dvc/machine/backend/terraform.py +++ b/dvc/machine/backend/terraform.py @@ -1,13 +1,10 @@ import logging import os -import tempfile from contextlib import contextmanager from typing import TYPE_CHECKING, Iterator, Optional from funcy import first -from python_terraform import Terraform -from dvc.exceptions import DvcException from dvc.fs.ssh import SSHFileSystem from .base import BaseMachineBackend @@ -18,51 +15,10 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -class TerraformError(DvcException): - pass - - -class DvcTerraform(Terraform): - def cmd(self, *args, **kwargs): - logger.debug(" ".join(self.generate_cmd_string(*args, **kwargs))) - kwargs["capture_output"] = False - ret, _stdout, _stderr = super().cmd(*args, **kwargs) - if ret != 0: - raise TerraformError("Cmd 'terraform {cmd}' failed") - - def iter_instances(self, name: str): - """Iterate over active iterative_machine instances.""" - self.read_state_file() - resources = getattr(self.tfstate, "resources", []) - for resource in resources: - if ( - resource.get("type") == "iterative_machine" - and resource.get("name") == name - ): - yield from ( - instance.get("attributes", {}) - for instance in resource.get("instances", []) - ) - - @staticmethod - @contextmanager - def pemfile(resource: dict): - """Generate a temporary SSH key (PEM) file for the specified resource. - - The path to the file will be yielded, and the file will be removed - upon exiting the yielded context. - """ - with tempfile.NamedTemporaryFile(delete=False) as fobj: - fobj.write(resource["ssh_private"].encode("ascii")) - try: - yield fobj.name - finally: - os.unlink(fobj.name) - - class TerraformBackend(BaseMachineBackend): @contextmanager def make_tf(self, name: str): + from dvc.tpi import DvcTerraform, TerraformError from dvc.utils.fs import makedirs try: @@ -110,6 +66,8 @@ class TerraformBackend(BaseMachineBackend): raise NotImplementedError def _default_resource(self, name): + from dvc.tpi import TerraformError + resource = first(self.instances(name)) if not resource: raise TerraformError(f"No active '{name}' instances") @@ -119,6 +77,8 @@ class TerraformBackend(BaseMachineBackend): def get_sshfs( self, name: Optional[str] = None, **config ) -> Iterator["SSHFileSystem"]: + from dvc.tpi import DvcTerraform + resource = self._default_resource(name) with DvcTerraform.pemfile(resource) as pem: fs = SSHFileSystem( @@ -129,6 +89,8 @@ class TerraformBackend(BaseMachineBackend): yield fs def run_shell(self, name: Optional[str] = None, **config): + from dvc.tpi import DvcTerraform + resource = self._default_resource(name) with DvcTerraform.pemfile(resource) as pem: self._shell( diff --git a/dvc/tpi/__init__.py b/dvc/tpi/__init__.py index de137038c..f39efd63f 100644 --- a/dvc/tpi/__init__.py +++ b/dvc/tpi/__init__.py @@ -1,9 +1,14 @@ import json import logging import os +import tempfile from collections import OrderedDict +from contextlib import contextmanager from jinja2 import Environment, FileSystemLoader, select_autoescape +from python_terraform import Terraform + +from dvc.exceptions import DvcException logger = logging.getLogger(__name__) @@ -30,6 +35,48 @@ _jinja_env = Environment( ) +class TerraformError(DvcException): + pass + + +class DvcTerraform(Terraform): + def cmd(self, *args, **kwargs): + logger.debug(" ".join(self.generate_cmd_string(*args, **kwargs))) + kwargs["capture_output"] = False + ret, _stdout, _stderr = super().cmd(*args, **kwargs) + if ret != 0: + raise TerraformError("Cmd 'terraform {cmd}' failed") + + def iter_instances(self, name: str): + """Iterate over active iterative_machine instances.""" + self.read_state_file() + resources = getattr(self.tfstate, "resources", []) + for resource in resources: + if ( + resource.get("type") == "iterative_machine" + and resource.get("name") == name + ): + yield from ( + instance.get("attributes", {}) + for instance in resource.get("instances", []) + ) + + @staticmethod + @contextmanager + def pemfile(resource: dict): + """Generate a temporary SSH key (PEM) file for the specified resource. + + The path to the file will be yielded, and the file will be removed + upon exiting the yielded context. + """ + with tempfile.NamedTemporaryFile(delete=False) as fobj: + fobj.write(resource["ssh_private"].encode("ascii")) + try: + yield fobj.name + finally: + os.unlink(fobj.name) + + def render_config(**config) -> str: """Render TPI iterative_machine config as terraform HCL.""" assert "name" in config and "cloud" in config diff --git a/tests/unit/machine/test_terraform.py b/tests/unit/machine/test_terraform.py index d45e1344f..8dc5bbd76 100644 --- a/tests/unit/machine/test_terraform.py +++ b/tests/unit/machine/test_terraform.py @@ -5,7 +5,7 @@ import pytest from mock import ANY, MagicMock from python_terraform import IsFlagged -from dvc.machine.backend.terraform import DvcTerraform, TerraformError +from dvc.tpi import DvcTerraform, TerraformError TEST_RESOURCE_STATE = """{ "version": 4,