Exemple #1
0
  def testAddKubectlNamespace(self):
    expected_tmpl = '/fake/kubectl --namespace {} get pods'

    # Checks that namespace args have been added.
    test_args = ['/fake/kubectl', 'get', 'pods']
    self.assertEqual(
        expected_tmpl.format(self.TEST_KUBECTL_DEFAULT_NAMESPACE).split(' '),
        command_util.AddKubectlNamespace(self.TEST_KUBECTL_DEFAULT_NAMESPACE,
                                         test_args))

    # Checks that a supplied namespace scope is not overwritten.
    test_args = ['/fake/kubectl', '--namespace', 'foo', 'get', 'pods']
    self.assertEqual(
        expected_tmpl.format('foo').split(' '),
        command_util.AddKubectlNamespace(self.TEST_KUBECTL_DEFAULT_NAMESPACE,
                                         test_args))
Exemple #2
0
        def _GetPodsCallback(args, **kwargs):
            get_pod_args = command_util.AddKubectlNamespace(
                kubectl_namespace, [self.TEST_KUBECTL_PATH, 'get', 'pods'])

            kubectl_util.AssertListHasPrefix(self, args, get_pod_args)
            if kwargs.get('out_func') is not None:
                kwargs['out_func'](output)
            return 0
    def Run(self, args):
        self.DeprecationWarningPrompt(args)
        self.CheckForRequiredCmdArgs(args)

        running_state = (api_util.GetMessagesModule(
            release_track=self.ReleaseTrack()).Environment.
                         StateValueValuesEnum.RUNNING)

        env_ref = args.CONCEPTS.environment.Parse()
        env_obj = environments_api_util.Get(env_ref,
                                            release_track=self.ReleaseTrack())

        if env_obj.state != running_state:
            raise command_util.Error(
                'Cannot execute subcommand for environment in state {}. '
                'Must be RUNNING.'.format(env_obj.state))

        cluster_id = env_obj.config.gkeCluster
        cluster_location_id = command_util.ExtractGkeClusterLocationId(env_obj)

        tty = 'no-tty' not in args

        with command_util.TemporaryKubeconfig(cluster_location_id, cluster_id):
            try:
                image_version = env_obj.config.softwareConfig.imageVersion
                airflow_version = self._ExtractAirflowVersion(image_version)

                self.CheckSubcommandAirflowSupport(args, airflow_version)
                self.CheckSubcommandNestedAirflowSupport(args, airflow_version)

                kubectl_ns = command_util.FetchKubectlNamespace(image_version)
                pod = command_util.GetGkePod(pod_substr=WORKER_POD_SUBSTR,
                                             kubectl_namespace=kubectl_ns)

                log.status.Print(
                    'Executing within the following Kubernetes cluster namespace: '
                    '{}'.format(kubectl_ns))

                self.BypassConfirmationPrompt(args, airflow_version)
                kubectl_args = ['exec', pod, '--stdin']
                if tty:
                    kubectl_args.append('--tty')
                kubectl_args.extend([
                    '--container', WORKER_CONTAINER, '--', 'airflow',
                    args.subcommand
                ])
                if args.subcommand_nested:
                    kubectl_args.append(args.subcommand_nested)
                if args.cmd_args:
                    kubectl_args.extend(args.cmd_args)

                command_util.RunKubectlCommand(
                    command_util.AddKubectlNamespace(kubectl_ns, kubectl_args),
                    out_func=log.out.Print)
            except command_util.KubectlError as e:
                raise self.ConvertKubectlError(e, env_obj)
Exemple #4
0
 def _KubectlExecCallback(args, **_):
     expected_args = command_util.AddKubectlNamespace(
         self.TEST_KUBECTL_DEFAULT_NAMESPACE, [
             self.TEST_KUBECTL_PATH, 'exec', pod, '--stdin', '--tty',
             '--container', self.TEST_CONTAINER, '--', 'airflow', subcmd
         ])
     if subcmd_args:
         expected_args.extend(subcmd_args)
     self.assertEqual(expected_args, args)
     return 0
Exemple #5
0
  def testRunKubectlCommand_Success(self, exec_mock,
                                    tmp_kubeconfig_mock):
    kubectl_args = ['exec', '-it', 'airflow-worker12345', 'bash']
    expected_args = command_util.AddKubectlNamespace(
        self.TEST_KUBECTL_DEFAULT_NAMESPACE,
        [self.TEST_KUBECTL_PATH] + kubectl_args)

    fake_exec = kubectl_util.FakeExec()
    exec_mock.side_effect = fake_exec
    tmp_kubeconfig_mock.side_effect = self.FakeTemporaryKubeconfig

    fake_exec.AddCallback(
        0, lambda args, **_: self.assertEqual(expected_args, args))
    command_util.RunKubectlCommand(
        kubectl_args, namespace=self.TEST_KUBECTL_DEFAULT_NAMESPACE)
    fake_exec.Verify()
  def Run(self, args):
    self.DeprecationWarningPrompt(args)

    running_state = (
        api_util.GetMessagesModule(release_track=self.ReleaseTrack())
        .Environment.StateValueValuesEnum.RUNNING)

    env_ref = args.CONCEPTS.environment.Parse()
    env_obj = environments_api_util.Get(
        env_ref, release_track=self.ReleaseTrack())

    if env_obj.state != running_state:
      raise command_util.Error(
          'Cannot execute subcommand for environment in state {}. '
          'Must be RUNNING.'.format(env_obj.state))

    cluster_id = env_obj.config.gkeCluster
    cluster_location_id = command_util.ExtractGkeClusterLocationId(env_obj)

    with command_util.TemporaryKubeconfig(cluster_location_id, cluster_id):
      try:
        kubectl_ns = command_util.FetchKubectlNamespace(
            env_obj.config.softwareConfig.imageVersion)
        pod = command_util.GetGkePod(
            pod_substr=WORKER_POD_SUBSTR, kubectl_namespace=kubectl_ns)

        log.status.Print(
            'Executing within the following kubectl namespace: {}'.format(
                kubectl_ns))

        self.BypassConfirmationPrompt(args)
        kubectl_args = [
            'exec', pod, '-tic', WORKER_CONTAINER, 'airflow', args.subcommand
        ]
        if args.cmd_args:
          # Add '--' to the argument list so kubectl won't eat the command args.
          kubectl_args.extend(['--'] + args.cmd_args)

        command_util.RunKubectlCommand(
            command_util.AddKubectlNamespace(kubectl_ns, kubectl_args),
            out_func=log.status.Print)
      except command_util.KubectlError as e:
        raise self.ConvertKubectlError(e, env_obj)
Exemple #7
0
    def Run(self, args):
        env_ref = args.CONCEPTS.environment.Parse()
        env_obj = environments_api_util.Get(env_ref,
                                            release_track=self.ReleaseTrack())

        cluster_id = env_obj.config.gkeCluster
        cluster_location_id = command_util.ExtractGkeClusterLocationId(env_obj)

        tty = 'no-tty' not in args

        with command_util.TemporaryKubeconfig(cluster_location_id, cluster_id):
            try:
                image_version = env_obj.config.softwareConfig.imageVersion

                kubectl_ns = command_util.FetchKubectlNamespace(image_version)
                pod = command_util.GetGkePod(pod_substr=WORKER_POD_SUBSTR,
                                             kubectl_namespace=kubectl_ns)

                log.status.Print(
                    'Executing within the following Kubernetes cluster namespace: '
                    '{}'.format(kubectl_ns))

                kubectl_args = ['exec', pod, '--stdin']
                if tty:
                    kubectl_args.append('--tty')
                kubectl_args.extend(['--container', WORKER_CONTAINER, '--'])
                if args.tree:
                    kubectl_args.extend(
                        ['python', '-m', 'pipdeptree', '--warn'])
                else:
                    kubectl_args.extend(['python', '-m', 'pip', 'list'])

                command_util.RunKubectlCommand(
                    command_util.AddKubectlNamespace(kubectl_ns, kubectl_args),
                    out_func=log.out.Print)
            except command_util.KubectlError as e:
                raise self.ConvertKubectlError(e, env_obj)