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)
