A contrib package for building AWS Lambdas from python code. (pantsbu…

…ild#6881)

Builds a pex, then runs lambdex to turn that pex into a bundle suitable for uploading to AWS Lambda.

contrib/awslambda/python/src/python/pants/__init__.py

@@ -0,0 +1 @@
+__import__('pkg_resources').declare_namespace(__name__)

contrib/awslambda/python/src/python/pants/contrib/__init__.py

@@ -0,0 +1 @@
+__import__('pkg_resources').declare_namespace(__name__)

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/BUILD

@@ -0,0 +1,17 @@
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+contrib_plugin(
+  name='plugin',
+  dependencies=[
+    'contrib/awslambda/python/src/python/pants/contrib/awslambda/python/subsystems',
+    'contrib/awslambda/python/src/python/pants/contrib/awslambda/python/targets',
+    'contrib/awslambda/python/src/python/pants/contrib/awslambda/python/tasks',
+    'src/python/pants/build_graph',
+    'src/python/pants/goal:task_registrar',
+  ],
+  distribution_name='pantsbuild.pants.contrib.awslambda',
+  description='AWS Lambda pants plugin.',
+  build_file_aliases=True,
+  register_goals=True,
+)

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/examples/BUILD

@@ -0,0 +1,29 @@
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+python_requirement_library(
+  name = 'pycountry',
+  requirements = [
+    python_requirement('pycountry==18.5.20'),
+  ],
+)
+
+python_library(
+  name='hello-lib',
+  sources = ['hello_lib.py'],
+)
+
+python_binary(
+  name='hello-bin',
+  source='hello_handler.py',
+  dependencies=[
+    ':pycountry',
+    ':hello-lib',
+  ]
+)
+
+python_awslambda(
+  name='hello-lambda',
+  binary=':hello-bin',
+  handler='pants.contrib.awslambda.python.examples.hello_handler:handler'
+)

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/examples/hello_handler.py

@@ -0,0 +1,14 @@
+# coding=utf-8
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import pycountry
+
+from pants.contrib.awslambda.python.examples.hello_lib import say_hello
+
+
+def handler(event, context):
+  usa = pycountry.countries.get(alpha_2='US').name
+  say_hello('from the {}'.format(usa))

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/examples/hello_lib.py

@@ -0,0 +1,9 @@
+# coding=utf-8
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+
+def say_hello(s):
+  print('Hello {}!'.format(s))

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/register.py

@@ -0,0 +1,25 @@
+# coding=utf-8
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from pants.build_graph.build_file_aliases import BuildFileAliases
+from pants.goal.task_registrar import TaskRegistrar as task
+
+from pants.contrib.awslambda.python.targets.python_awslambda import PythonAWSLambda
+from pants.contrib.awslambda.python.tasks.lambdex_prep import LambdexPrep
+from pants.contrib.awslambda.python.tasks.lambdex_run import LambdexRun
+
+
+def build_file_aliases():
+  return BuildFileAliases(
+    targets={
+      'python_awslambda': PythonAWSLambda,
+    }
+  )
+
+
+def register_goals():
+  task(name='lambdex-prep', action=LambdexPrep).install('bundle')
+  task(name='lambdex-run', action=LambdexRun).install('bundle')

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/subsystems/BUILD

@@ -0,0 +1,8 @@
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+python_library(
+  dependencies = [
+    'src/python/pants/backend/python/subsystems',
+  ]
+)

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/subsystems/lambdex.py

@@ -0,0 +1,13 @@
+# coding=utf-8
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from pants.backend.python.subsystems.python_tool_base import PythonToolBase
+
+
+class Lambdex(PythonToolBase):
+  options_scope = 'lambdex'
+  default_requirements = ['lambdex==0.1.2']
+  default_entry_point = 'lambdex.bin.lambdex'

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/targets/BUILD

@@ -0,0 +1,8 @@
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+python_library(
+  dependencies = [
+    'src/python/pants/backend/python/targets',
+  ]
+)

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/targets/python_awslambda.py

@@ -0,0 +1,64 @@
+# coding=utf-8
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from pants.backend.python.targets.python_binary import PythonBinary
+from pants.base.exceptions import TargetDefinitionException
+from pants.base.payload import Payload
+from pants.base.payload_field import PrimitiveField
+from pants.build_graph.target import Target
+
+
+class PythonAWSLambda(Target):
+  """A self-contained Python function suitable for uploading to AWS Lambda.
+
+  :API: public
+  """
+
+  def __init__(self,
+               binary=None,
+               handler=None,
+               **kwargs):
+    """
+    :param string binary: Target spec of the ``python_binary`` that contains the handler.
+    :param string handler: Lambda handler entrypoint (module.dotted.name:handler_func).
+    """
+    payload = Payload()
+    payload.add_fields({
+      'binary': PrimitiveField(binary),
+      'handler': PrimitiveField(handler),
+    })
+    super(PythonAWSLambda, self).__init__(payload=payload, **kwargs)
+
+  @classmethod
+  def alias(cls):
+    return 'python_awslambda'
+
+  @classmethod
+  def compute_dependency_specs(cls, kwargs=None, payload=None):
+    for spec in super(PythonAWSLambda, cls).compute_dependency_specs(kwargs, payload):
+      yield spec
+    target_representation = kwargs or payload.as_dict()
+    binary = target_representation.get('binary')
+    if binary:
+      yield binary
+
+  @property
+  def binary(self):
+    """Returns the binary that builds the pex for this lambda."""
+    dependencies = self.dependencies
+    if len(dependencies) != 1:
+      raise TargetDefinitionException(self, 'An app must define exactly one binary '
+                                            'dependency, have: {}'.format(dependencies))
+    binary = dependencies[0]
+    if not isinstance(binary, PythonBinary):
+      raise TargetDefinitionException(self, 'Expected binary dependency to be a python_binary '
+                                            'target, found {}'.format(binary))
+    return binary
+
+  @property
+  def handler(self):
+    """Return the handler function for the lambda."""
+    return self.payload.handler

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/tasks/BUILD

@@ -0,0 +1,9 @@
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+python_library(
+  dependencies = [
+    'contrib/awslambda/python/src/python/pants/contrib/awslambda/python/subsystems',
+    'src/python/pants/backend/python/tasks',
+  ]
+)

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/tasks/lambdex_prep.py

@@ -0,0 +1,18 @@
+# coding=utf-8
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from pants.backend.python.tasks.python_tool_prep_base import PythonToolInstance, PythonToolPrepBase
+
+from pants.contrib.awslambda.python.subsystems.lambdex import Lambdex
+
+
+class LambdexInstance(PythonToolInstance):
+  pass
+
+
+class LambdexPrep(PythonToolPrepBase):
+  tool_subsystem_cls = Lambdex
+  tool_instance_cls = LambdexInstance

contrib/awslambda/python/src/python/pants/contrib/awslambda/python/tasks/lambdex_run.py

@@ -0,0 +1,99 @@
+# coding=utf-8
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import functools
+import os
+import shutil
+
+from pants.base.build_environment import get_buildroot
+from pants.base.exceptions import TaskError
+from pants.base.workunit import WorkUnitLabel
+from pants.task.task import Task
+from pants.util.contextutil import temporary_dir
+from pants.util.dirutil import safe_mkdir_for
+from pants.util.fileutil import atomic_copy
+
+from pants.contrib.awslambda.python.targets.python_awslambda import PythonAWSLambda
+from pants.contrib.awslambda.python.tasks.lambdex_prep import LambdexPrep
+
+
+class LambdexRun(Task):
+  """Runs Lambdex to convert pexes to AWS Lambda functions.
+
+  Note that your pex must be built to run on Amazon Linux, e.g., if it contains native code.
+  When deploying the lambda, its handler should be set to `lambdex_handler.handler`, which
+  is a wrapper around the target-specified handler.
+  """
+
+  @staticmethod
+  def _is_python_lambda(target):
+    return isinstance(target, PythonAWSLambda)
+
+  @classmethod
+  def product_types(cls):
+    return ['python_aws_lambda']
+
+  @classmethod
+  def prepare(cls, options, round_manager):
+    super(LambdexRun, cls).prepare(options, round_manager)
+    round_manager.require_data(LambdexPrep.tool_instance_cls)
+    round_manager.require('pex_archives')
+
+  @classmethod
+  def create_target_dirs(self):
+    return True
+
+  def execute(self):
+    targets = self.get_targets(self._is_python_lambda)
+    with self.invalidated(targets=targets) as invalidation_check:
+      python_lambda_product = self.context.products.get_data('python_aws_lambda', dict)
+      for vt in invalidation_check.all_vts:
+        lambda_path = os.path.join(vt.results_dir, '{}.pex'.format(vt.target.name))
+        if not vt.valid:
+          self.context.log.debug('Existing lambda for {} is invalid, rebuilding'.format(vt.target))
+          self._create_lambda(vt.target, lambda_path)
+        else:
+          self.context.log.debug('Using existing lambda for {}'.format(vt.target))
+
+        python_lambda_product[vt.target] = lambda_path
+        self.context.log.debug('created {}'.format(os.path.relpath(lambda_path, get_buildroot())))
+
+        # Put a copy in distdir.
+        lambda_copy = os.path.join(self.get_options().pants_distdir, os.path.basename(lambda_path))
+        safe_mkdir_for(lambda_copy)
+        atomic_copy(lambda_path, lambda_copy)
+        self.context.log.info('created lambda {}'.format(
+          os.path.relpath(lambda_copy, get_buildroot())))
+
+  def _create_lambda(self, target, lambda_path):
+    orig_pex_path = self._get_pex_path(target.binary)
+    with temporary_dir() as tmpdir:
+      # lambdex modifies the pex in-place, so we operate on a copy.
+      tmp_lambda_path = os.path.join(tmpdir, os.path.basename(lambda_path))
+      shutil.copy(orig_pex_path, tmp_lambda_path)
+      lambdex = self.context.products.get_data(LambdexPrep.tool_instance_cls)
+      workunit_factory = functools.partial(self.context.new_workunit,
+                                           name='run-lambdex',
+                                           labels=[WorkUnitLabel.TOOL, WorkUnitLabel.TOOL])
+      cmdline, exit_code = lambdex.run(workunit_factory,
+                                       ['build', '-e', target.handler, tmp_lambda_path])
+      if exit_code != 0:
+        raise TaskError('{} ... exited non-zero ({}).'.format(cmdline, exit_code),
+                        exit_code=exit_code)
+      shutil.move(tmp_lambda_path, lambda_path)
+    return lambda_path
+
+  # TODO(benjy): Switch python_binary_create to use data products, and get rid of this wrinkle
+  # here and in python_bundle.py.
+  def _get_pex_path(self, binary_tgt):
+    pex_archives = self.context.products.get('pex_archives')
+    paths = []
+    for basedir, filenames in pex_archives.get(binary_tgt).items():
+      for filename in filenames:
+        paths.append(os.path.join(basedir, filename))
+    if len(paths) != 1:
+      raise TaskError('Expected one binary but found: {}'.format(', '.join(sorted(paths))))
+    return paths[0]

contrib/awslambda/python/tests/python/pants_test/__init__.py

@@ -0,0 +1 @@
+__import__('pkg_resources').declare_namespace(__name__)

contrib/awslambda/python/tests/python/pants_test/contrib/__init__.py

@@ -0,0 +1 @@
+__import__('pkg_resources').declare_namespace(__name__)

contrib/awslambda/python/tests/python/pants_test/contrib/awslambda/python/BUILD

@@ -0,0 +1,12 @@
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+python_tests(
+  dependencies=[
+    'contrib/awslambda/python/src/python/pants/contrib/awslambda/python:plugin',
+    'src/python/pants/util:contextutil',
+    'src/python/pants/util:process_handler',
+    'tests/python/pants_test:int-test',
+  ],
+  tags={'integration'}
+)

contrib/awslambda/python/tests/python/pants_test/contrib/awslambda/python/__init__.py

@@ -0,0 +1 @@
+__import__('pkg_resources').declare_namespace(__name__)

contrib/awslambda/python/tests/python/pants_test/contrib/awslambda/python/test_python_awslambda_integration.py

@@ -0,0 +1,43 @@
+# coding=utf-8
+# Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import os
+
+from pants.util.contextutil import temporary_dir
+from pants.util.process_handler import subprocess
+from pants_test.pants_run_integration_test import PantsRunIntegrationTest
+
+
+class PythonAWSLambdaIntegrationTest(PantsRunIntegrationTest):
+  @classmethod
+  def hermetic(cls):
+    return True
+
+  def test_awslambda_bundle(self):
+    with temporary_dir() as distdir:
+      config = {
+        'GLOBAL': {
+          'pants_distdir': distdir,
+          'pythonpath': ['%(buildroot)s/contrib/awslambda/python/src/python'],
+          'backend_packages': ['pants.backend.python', 'pants.contrib.awslambda.python'],
+        }
+      }
+
+      command = [
+        'bundle',
+        'contrib/awslambda/python/src/python/pants/contrib/awslambda/python/examples:hello-lambda',
+      ]
+      pants_run = self.run_pants(command=command, config=config)
+      self.assert_success(pants_run)
+
+      # Now run the lambda via the wrapper handler injected by lambdex (note that this
+      # is distinct from the pex's entry point - a handler must be a function with two arguments,
+      # whereas the pex entry point is a module).
+      awslambda = os.path.join(distdir, 'hello-lambda.pex')
+      output = subprocess.check_output(env={'PEX_INTERPRETER': '1'}, args=[
+        '{} -c "from lambdex_handler import handler; handler(None, None)"'.format(awslambda)
+      ], shell=True)
+      self.assertEquals(b'Hello from the United States!', output.strip())