def _ConstructLabelsPatch(clear_labels,
remove_labels,
update_labels,
release_track=base.ReleaseTrack.GA):
"""Constructs an environment patch for updating labels.
Args:
clear_labels: bool, whether to clear the labels dictionary.
remove_labels: iterable(string), Iterable of label names to remove.
update_labels: {string: string}, dict of label names and values to set.
release_track: base.ReleaseTrack, the release track of command. Will dictate
which Composer client library will be used.
Returns:
(str, Environment), the field mask and environment to use for update.
"""
messages = api_util.GetMessagesModule(release_track=release_track)
env_cls = messages.Environment
entry_cls = env_cls.LabelsValue.AdditionalProperty
def _BuildEnv(entries):
return env_cls(labels=env_cls.LabelsValue(
additionalProperties=entries))
return command_util.BuildPartialUpdate(clear_labels, remove_labels,
update_labels, 'labels', entry_cls,
_BuildEnv)
def _ConstructPyPiPackagesPatch(clear_pypi_packages,
remove_pypi_packages,
update_pypi_packages,
release_track=base.ReleaseTrack.GA):
"""Constructs an environment patch for partially updating PyPI packages.
Args:
clear_pypi_packages: bool, whether to clear the PyPI packages dictionary.
remove_pypi_packages: iterable(string), Iterable of PyPI package names to
remove.
update_pypi_packages: {string: string}, dict mapping PyPI package name
to optional extras and version specifier.
release_track: base.ReleaseTrack, the release track of command. Will dictate
which Composer client library will be used.
Returns:
(str, Environment), the field mask and environment to use for update.
"""
messages = api_util.GetMessagesModule(release_track=release_track)
env_cls = messages.Environment
pypi_packages_cls = (messages.SoftwareConfig.PypiPackagesValue)
entry_cls = pypi_packages_cls.AdditionalProperty
def _BuildEnv(entries):
software_config = messages.SoftwareConfig(
pypiPackages=pypi_packages_cls(additionalProperties=entries))
config = messages.EnvironmentConfig(softwareConfig=software_config)
return env_cls(config=config)
return command_util.BuildPartialUpdate(
clear_pypi_packages, remove_pypi_packages, update_pypi_packages,
'config.software_config.pypi_packages', entry_cls, _BuildEnv)
def _ConstructAirflowConfigsPatch(clear_airflow_configs,
remove_airflow_configs,
update_airflow_configs,
release_track=base.ReleaseTrack.GA):
"""Constructs an environment patch for updating Airflow configs.
Args:
clear_airflow_configs: bool, whether to clear the Airflow configs
dictionary.
remove_airflow_configs: iterable(string), Iterable of Airflow config
property names to remove.
update_airflow_configs: {string: string}, dict of Airflow config property
names and values to set.
release_track: base.ReleaseTrack, the release track of command. Will dictate
which Composer client library will be used.
Returns:
(str, Environment), the field mask and environment to use for update.
"""
messages = api_util.GetMessagesModule(release_track=release_track)
env_cls = messages.Environment
airflow_config_overrides_cls = (
messages.SoftwareConfig.AirflowConfigOverridesValue)
entry_cls = airflow_config_overrides_cls.AdditionalProperty
def _BuildEnv(entries):
software_config = messages.SoftwareConfig(
airflowConfigOverrides=airflow_config_overrides_cls(
additionalProperties=entries))
config = messages.EnvironmentConfig(softwareConfig=software_config)
return env_cls(config=config)
return command_util.BuildPartialUpdate(
clear_airflow_configs, remove_airflow_configs, update_airflow_configs,
'config.software_config.airflow_config_overrides', entry_cls, _BuildEnv)
def _BuildPartialUpdateTestHelper(self,
expected_field_mask,
expected_patch,
clear,
remove_keys=None,
set_entries=None):
"""Helper for testing BuildPartialUpdate.
Using this helper, one can test that BuildPartialUpdate combines `clear`,
`remove_keys`, and `set_entries` (in that order) to build a patch object and
update mask.
Args:
expected_field_mask: str, sorted, comma-delimited list of fields expected
to appear in the patch request's field mask
expected_patch: dict, a dictionary of the key-value pairs expected to
be in the patch
clear: bool, the value of the clear parameter to pass through
remove_keys: [str], the value of the remove_keys parameter to pass through
set_entries: {str: str}, the value of the set_entries parameter to pass
through
"""
actual_field_mask, actual_patch = command_util.BuildPartialUpdate(
clear=clear,
remove_keys=remove_keys,
set_entries=set_entries,
field_mask_prefix=self.field_mask_prefix,
entry_cls=UtilGATest._FakeEntry,
env_builder=UtilGATest._FakePatchBuilder)
self.assertEqual(expected_field_mask, actual_field_mask)
self.assertEqual(expected_patch, actual_patch)
def _ConstructLabelsPatch(clear_labels, remove_labels, update_labels):
"""Constructs an environment patch for updating labels.
Args:
clear_labels: bool, whether to clear the labels dictionary.
remove_labels: iterable(string), Iterable of label names to remove.
update_labels: {string: string}, dict of label names and values to set.
Returns:
(str, Environment), the field mask and environment to use for update.
"""
messages = api_util.GetMessagesModule()
env_cls = messages.Environment
entry_cls = env_cls.LabelsValue.AdditionalProperty
def _BuildEnv(entries):
return env_cls(labels=env_cls.LabelsValue(
additionalProperties=entries))
return command_util.BuildPartialUpdate(clear_labels, remove_labels,
update_labels, 'labels', entry_cls,
_BuildEnv)
