def create_cluster(
self,
cluster: Union[Dict, Cluster, None],
project_id: str = PROVIDE_PROJECT_ID,
retry: Union[Retry, _MethodDefault] = DEFAULT,
timeout: Optional[float] = None,
) -> str:
"""
Creates a cluster, consisting of the specified number and type of Google Compute
Engine instances.
:param cluster: A Cluster protobuf or dict. If dict is provided, it must
be of the same form as the protobuf message
:class:`google.cloud.container_v1.types.Cluster`
:param project_id: Google Cloud project ID
:param retry: A retry object (``google.api_core.retry.Retry``) used to
retry requests.
If None is specified, requests will not be retried.
:param timeout: The amount of time, in seconds, to wait for the request to
complete. Note that if retry is specified, the timeout applies to each
individual attempt.
:return: The full url to the new, or existing, cluster
:raises:
ParseError: On JSON parsing problems when trying to convert dict
AirflowException: cluster is not dict type nor Cluster proto type
"""
if isinstance(cluster, dict):
cluster = Cluster.from_json(json.dumps(cluster))
elif not isinstance(cluster, Cluster):
raise AirflowException(
"cluster is not instance of Cluster proto or python dict")
self._append_label(cluster, 'airflow-version',
'v' + version.version) # type: ignore
self.log.info(
"Creating (project_id=%s, location=%s, cluster_name=%s)",
project_id,
self.location,
cluster.name, # type: ignore
)
try:
resource = self.get_cluster_manager_client().create_cluster(
parent=f'projects/{project_id}/locations/{self.location}',
cluster=cluster, # type: ignore
retry=retry,
timeout=timeout,
)
resource = self.wait_for_operation(resource)
return resource.target_link
except AlreadyExists as error:
self.log.info('Assuming Success: %s', error.message)
return self.get_cluster(name=cluster.name,
project_id=project_id) # type: ignore
def persist(context: "Context", task_instance, cluster: Union[Dict, Cluster, None]):
if isinstance(cluster, dict):
cluster = Cluster.from_json(json.dumps(cluster))
task_instance.xcom_push(
context=context,
key=KubernetesEngineClusterLink.key,
value={
"location": task_instance.location,
"cluster_name": cluster.name, # type: ignore
"project_id": task_instance.project_id,
},
)
