def test_create_instance_when_exists(self, mock_get_conn): mock_get_conn.return_value.get_instance.return_value = Instance(name=TEST_NAME) result = self.hook.create_instance( location=TEST_LOCATION, instance_id=TEST_INSTANCE_ID, instance=Instance(name=TEST_NAME), project_id=TEST_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) mock_get_conn.return_value.get_instance.assert_called_once_with( name="projects/test-project-id/locations/test-location/instances/test-instance-id", retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) self.assertEqual(Instance(name=TEST_NAME), result)
def test_update_instance_without_project_id(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): self.hook.update_instance( # pylint: disable=no-value-for-parameter update_mask=TEST_UPDATE_MASK, instance=Instance(name=TEST_NAME), retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, )
def test_create_instance_when_not_exists(self, mock_get_conn, mock_project_id): mock_get_conn.return_value.get_instance.side_effect = [ NotFound("Instance not found"), Instance(name=TEST_NAME), ] mock_get_conn.return_value.create_instance.return_value.result.return_value = Instance( name=TEST_NAME) result = self.hook.create_instance( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, instance_id=TEST_INSTANCE_ID, instance=Instance(name=TEST_NAME), retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) mock_get_conn.return_value.get_instance.has_calls([ mock.call(name=TEST_NAME, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA), mock.call(name=TEST_NAME, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA), ]) mock_get_conn.return_value.create_instance.assert_called_once_with( request=dict( parent=TEST_PARENT_DEFAULT_PROJECT_ID, instance=Instance( name=TEST_NAME, labels={ "airflow-version": "v" + version.version.replace(".", "-").replace("+", "-") }, ), instance_id=TEST_INSTANCE_ID, ), metadata=TEST_METADATA, retry=TEST_RETRY, timeout=TEST_TIMEOUT, ) assert Instance(name=TEST_NAME) == result
def test_create_instance_when_not_exists(self, mock_get_conn): mock_get_conn.return_value.get_instance.side_effect = [ NotFound("Instance not found"), Instance(name=TEST_NAME), ] mock_get_conn.return_value.create_instance.return_value.result.return_value = Instance(name=TEST_NAME) result = self.hook.create_instance( location=TEST_LOCATION, instance_id=TEST_INSTANCE_ID, instance=Instance(name=TEST_NAME), project_id=TEST_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) mock_get_conn.return_value.get_instance.has_calls( [ mock.call( name="projects/test-project-id/locations/test-location/instances/test-instance-id", retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ), mock.call( name="projects/test-project-id/locations/test-location/instances/test-instance-id", retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ), ] ) mock_get_conn.return_value.create_instance.assert_called_once_with( instance=Instance( name=TEST_NAME, labels={"airflow-version": "v" + version.version.replace(".", "-").replace("+", "-")}, ), instance_id=TEST_INSTANCE_ID, metadata=TEST_METADATA, parent=TEST_PARENT, retry=TEST_RETRY, timeout=TEST_TIMEOUT, ) self.assertEqual(Instance(name=TEST_NAME), result)
def test_get_instance_without_project_id(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): self.hook.get_instance( location=TEST_LOCATION, instance=Instance(name=TEST_NAME), project_id=None, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, )
CloudMemorystoreCreateInstanceOperator, CloudMemorystoreDeleteInstanceOperator, CloudMemorystoreExportInstanceOperator, CloudMemorystoreFailoverInstanceOperator, CloudMemorystoreGetInstanceOperator, CloudMemorystoreImportOperator, CloudMemorystoreListInstancesOperator, CloudMemorystoreScaleInstanceOperator, CloudMemorystoreUpdateInstanceOperator, ) TEST_GCP_CONN_ID = "test-gcp-conn-id" TEST_TASK_ID = "task-id" TEST_LOCATION = "test-location" TEST_INSTANCE_ID = "test-instance-id" TEST_INSTANCE = Instance(name="instance") TEST_INSTANCE_NAME = "test-instance-name" TEST_PROJECT_ID = "test-project-id" TEST_RETRY = Retry() # type: Retry TEST_TIMEOUT = 10 # type: float TEST_INSTANCE_SIZE = 4 # type: int TEST_METADATA = [("KEY", "VALUE")] # type: Sequence[Tuple[str, str]] TEST_OUTPUT_CONFIG = { "gcs_destination": { "uri": "gs://test-bucket/file.rdb" } } # type: Dict TEST_DATA_PROTECTION_MODE = FailoverInstanceRequest.DataProtectionMode.LIMITED_DATA_LOSS TEST_INPUT_CONFIG = { "gcs_source": { "uri": "gs://test-bucket/file.rdb"
def create_instance( self, location: str, instance_id: str, instance: Union[Dict, Instance], project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ): """ Creates a Redis instance based on the specified tier and memory size. By default, the instance is accessible from the project's `default network <https://cloud.google.com/compute/docs/networks-and-firewalls#networks>`__. :param location: The location of the Cloud Memorystore instance (for example europe-west1) :type location: str :param instance_id: Required. The logical name of the Redis instance in the customer project with the following restrictions: - Must contain only lowercase letters, numbers, and hyphens. - Must start with a letter. - Must be between 1-40 characters. - Must end with a number or a letter. - Must be unique within the customer project / location :type instance_id: str :param instance: Required. A Redis [Instance] resource If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.redis_v1.types.Instance` :type instance: Union[Dict, google.cloud.redis_v1.types.Instance] :param project_id: Project ID of the project that contains the instance. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :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. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ client = self.get_conn() parent = CloudRedisClient.location_path(project_id, location) instance_name = CloudRedisClient.instance_path(project_id, location, instance_id) try: instance = client.get_instance(name=instance_name, retry=retry, timeout=timeout, metadata=metadata) self.log.info("Instance exists. Skipping creation.") return instance except NotFound: self.log.info("Instance not exists.") if isinstance(instance, dict): instance = ParseDict(instance, Instance()) elif not isinstance(instance, Instance): raise AirflowException( "instance is not instance of Instance type or python dict") self._append_label(instance, "airflow-version", "v" + version.version) result = client.create_instance( parent=parent, instance_id=instance_id, instance=instance, retry=retry, timeout=timeout, metadata=metadata, ) result.result() self.log.info("Instance created.") return client.get_instance(name=instance_name, retry=retry, timeout=timeout, metadata=metadata)
def update_instance( self, update_mask: Union[Dict, FieldMask], instance: Union[Dict, Instance], location: Optional[str] = None, instance_id: Optional[str] = None, project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ): """ Updates the metadata and configuration of a specific Redis instance. :param update_mask: Required. Mask of fields to update. At least one path must be supplied in this field. The elements of the repeated paths field may only include these fields from ``Instance``: - ``displayName`` - ``labels`` - ``memorySizeGb`` - ``redisConfig`` If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.redis_v1.types.FieldMask` :type update_mask: Union[Dict, google.cloud.redis_v1.types.FieldMask] :param instance: Required. Update description. Only fields specified in ``update_mask`` are updated. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.redis_v1.types.Instance` :type instance: Union[Dict, google.cloud.redis_v1.types.Instance] :param location: The location of the Cloud Memorystore instance (for example europe-west1) :type location: str :param instance_id: The logical name of the Redis instance in the customer project. :type instance_id: str :param project_id: Project ID of the project that contains the instance. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :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. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ client = self.get_conn() if isinstance(instance, dict): instance = ParseDict(instance, Instance()) elif not isinstance(instance, Instance): raise AirflowException( "instance is not instance of Instance type or python dict") if location and instance_id: name = CloudRedisClient.instance_path(project_id, location, instance_id) instance.name = name self.log.info("Updating instances: %s", instance.name) result = client.update_instance(update_mask=update_mask, instance=instance, retry=retry, timeout=timeout, metadata=metadata) result.result() self.log.info("Instance updated: %s", instance.name)