def __init__(
self,
*,
resource_id: str,
body_patch: dict,
project_id: Optional[str] = None,
request_id: Optional[str] = None,
gcp_conn_id: str = 'google_cloud_default',
api_version: str = 'v1',
validate_body: bool = True,
impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
**kwargs,
) -> None:
self.body_patch = body_patch
self.request_id = request_id
self._field_validator = None # Optional[GcpBodyFieldValidator]
if 'name' not in self.body_patch:
raise AirflowException(
f"The body '{body_patch}' should contain at least name for the new operator "
f"in the 'name' field")
if validate_body:
self._field_validator = GcpBodyFieldValidator(
GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION,
api_version=api_version)
self._field_sanitizer = GcpBodyFieldSanitizer(
GCE_INSTANCE_TEMPLATE_FIELDS_TO_SANITIZE)
super().__init__(
project_id=project_id,
zone='global',
resource_id=resource_id,
gcp_conn_id=gcp_conn_id,
api_version=api_version,
impersonation_chain=impersonation_chain,
**kwargs,
)
def test_sanitize_should_not_fail_if_field_is_absent_in_body(self):
body = {}
fields_to_sanitize = ["kind"]
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
sanitizer.sanitize(body)
self.assertEqual({}, body)
def test_sanitize_should_remove_for_multiple_fields_from_root_level(self):
body = {"kind": "compute#instanceTemplate", "name": "instance"}
fields_to_sanitize = ["kind", "name"]
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
sanitizer.sanitize(body)
self.assertEqual({}, body)
def test_sanitize_should_not_fail_with_none_body(self):
body = None
fields_to_sanitize = []
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
sanitizer.sanitize(body)
self.assertIsNone(body)
def test_sanitize_should_fail_with_none_fields(self):
body = {}
fields_to_sanitize = None
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
with self.assertRaises(TypeError):
sanitizer.sanitize(body)
def test_sanitize_should_sanitize_empty_body_and_fields(self):
body = {}
fields_to_sanitize = []
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
sanitizer.sanitize(body)
self.assertEqual({}, body)
def test_sanitize_should_remove_all_fields_in_any_nested_body(self):
fields_to_sanitize = [
"kind",
"properties.disks.kind",
"properties.metadata.kind",
]
body = {
"kind": "compute#instanceTemplate",
"name": "instance",
"properties": {
"disks": [
{
"name": "a",
"kind": "compute#attachedDisk",
"type": "PERSISTENT",
"mode": "READ_WRITE",
},
{
"name": "b",
"kind": "compute#attachedDisk",
"type": "PERSISTENT",
"mode": "READ_WRITE",
},
],
"metadata": {
"kind": "compute#metadata",
"fingerprint": "GDPUYxlwHe4="
},
},
}
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
sanitizer.sanitize(body)
self.assertEqual(
{
"name": "instance",
"properties": {
"disks": [
{
"name": "a",
"type": "PERSISTENT",
"mode": "READ_WRITE"
},
{
"name": "b",
"type": "PERSISTENT",
"mode": "READ_WRITE"
},
],
"metadata": {
"fingerprint": "GDPUYxlwHe4="
},
},
},
body,
)
def test_sanitize_should_not_remove_fields_for_incorrect_specification(self):
actual_body = [
{"kind": "compute#instanceTemplate", "name": "instance"},
{"kind": "compute#instanceTemplate1", "name": "instance1"},
{"kind": "compute#instanceTemplate2", "name": "instance2"},
]
body = deepcopy(actual_body)
fields_to_sanitize = ["kind"]
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
sanitizer.sanitize(body)
self.assertEqual(actual_body, body)
def test_sanitize_should_remove_all_fields_in_a_list_value(self):
body = {"fields": [
{"kind": "compute#instanceTemplate", "name": "instance"},
{"kind": "compute#instanceTemplate1", "name": "instance1"},
{"kind": "compute#instanceTemplate2", "name": "instance2"},
]}
fields_to_sanitize = ["fields.kind"]
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
sanitizer.sanitize(body)
self.assertEqual({"fields": [
{"name": "instance"},
{"name": "instance1"},
{"name": "instance2"},
]}, body)
def test_sanitize_should_not_fail_if_no_specification_matches(self):
fields_to_sanitize = [
"properties.disks.kind1",
"properties.metadata.kind2",
]
body = {"name": "instance", "properties": {"disks": None}}
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
sanitizer.sanitize(body)
self.assertEqual({
"name": "instance",
"properties": {
"disks": None
}
}, body)
def test_sanitize_should_not_fail_if_type_in_body_do_not_match_with_specification(
self):
fields_to_sanitize = [
"properties.disks.kind",
"properties.metadata.kind2",
]
body = {"name": "instance", "properties": {"disks": 1}}
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
sanitizer.sanitize(body)
self.assertEqual({
"name": "instance",
"properties": {
"disks": 1
}
}, body)
def test_sanitize_should_not_fail_if_specification_has_none_value(self):
fields_to_sanitize = [
"kind",
"properties.disks.kind",
"properties.metadata.kind",
]
body = {
"kind": "compute#instanceTemplate",
"name": "instance",
"properties": {
"disks": None
}
}
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
sanitizer.sanitize(body)
self.assertEqual({
"name": "instance",
"properties": {
"disks": None
}
}, body)
class ComputeEngineCopyInstanceTemplateOperator(ComputeEngineBaseOperator):
"""
Copies the instance template, applying specified changes.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:ComputeEngineCopyInstanceTemplateOperator`
:param resource_id: Name of the Instance Template
:type resource_id: str
:param body_patch: Patch to the body of instanceTemplates object following rfc7386
PATCH semantics. The body_patch content follows
https://cloud.google.com/compute/docs/reference/rest/v1/instanceTemplates
Name field is required as we need to rename the template,
all the other fields are optional. It is important to follow PATCH semantics
- arrays are replaced fully, so if you need to update an array you should
provide the whole target array as patch element.
:type body_patch: dict
:param project_id: Optional, Google Cloud Platform Project ID where the Compute
Engine Instance exists. If set to None or missing, the default project_id from the GCP connection
is used.
:type project_id: str
:param request_id: Optional, unique request_id that you might add to achieve
full idempotence (for example when client call times out repeating the request
with the same request id will not create a new instance template again).
It should be in UUID format as defined in RFC 4122.
:type request_id: str
:param gcp_conn_id: Optional, The connection ID used to connect to Google Cloud
Platform. Defaults to 'google_cloud_default'.
:type gcp_conn_id: str
:param api_version: Optional, API version used (for example v1 - or beta). Defaults
to v1.
:type api_version: str
:param validate_body: Optional, If set to False, body validation is not performed.
Defaults to False.
:type validate_body: bool
"""
# [START gce_instance_template_copy_operator_template_fields]
template_fields = ('project_id', 'resource_id', 'request_id',
'gcp_conn_id', 'api_version')
# [END gce_instance_template_copy_operator_template_fields]
@apply_defaults
def __init__(self, *,
resource_id: str,
body_patch: dict,
project_id: Optional[str] = None,
request_id: Optional[str] = None,
gcp_conn_id: str = 'google_cloud_default',
api_version: str = 'v1',
validate_body: bool = True,
**kwargs) -> None:
self.body_patch = body_patch
self.request_id = request_id
self._field_validator = None # Optional[GcpBodyFieldValidator]
if 'name' not in self.body_patch:
raise AirflowException("The body '{}' should contain at least "
"name for the new operator in the 'name' field".
format(body_patch))
if validate_body:
self._field_validator = GcpBodyFieldValidator(
GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION, api_version=api_version)
self._field_sanitizer = GcpBodyFieldSanitizer(
GCE_INSTANCE_TEMPLATE_FIELDS_TO_SANITIZE)
super().__init__(
project_id=project_id, zone='global', resource_id=resource_id,
gcp_conn_id=gcp_conn_id, api_version=api_version, **kwargs)
def _validate_all_body_fields(self):
if self._field_validator:
self._field_validator.validate(self.body_patch)
def execute(self, context):
hook = ComputeEngineHook(gcp_conn_id=self.gcp_conn_id, api_version=self.api_version)
self._validate_all_body_fields()
try:
# Idempotence check (sort of) - we want to check if the new template
# is already created and if is, then we assume it was created by previous run
# of CopyTemplate operator - we do not check if content of the template
# is as expected. Templates are immutable so we cannot update it anyway
# and deleting/recreating is not worth the hassle especially
# that we cannot delete template if it is already used in some Instance
# Group Manager. We assume success if the template is simply present
existing_template = hook.get_instance_template(
resource_id=self.body_patch['name'], project_id=self.project_id)
self.log.info(
"The %s template already existed. It was likely created by previous run of the operator. "
"Assuming success.",
existing_template
)
return existing_template
except HttpError as e:
# We actually expect to get 404 / Not Found here as the template should
# not yet exist
if not e.resp.status == 404:
raise e
old_body = hook.get_instance_template(resource_id=self.resource_id,
project_id=self.project_id)
new_body = deepcopy(old_body)
self._field_sanitizer.sanitize(new_body)
new_body = merge(new_body, self.body_patch)
self.log.info("Calling insert instance template with updated body: %s", new_body)
hook.insert_instance_template(body=new_body,
request_id=self.request_id,
project_id=self.project_id)
return hook.get_instance_template(resource_id=self.body_patch['name'],
project_id=self.project_id)