def sink_delete(self, project, sink_name):
"""API call: delete a sink resource.
:type project: str
:param project: ID of the project containing the sink.
:type sink_name: str
:param sink_name: the name of the sink
"""
options = None
path = 'projects/%s/sinks/%s' % (project, sink_name)
try:
self._gax_api.delete_sink(path, options=options)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(path)
raise
def metric_delete(self, project, metric_name):
"""API call: delete a metric resource.
:type project: str
:param project: ID of the project containing the metric.
:type metric_name: str
:param metric_name: the name of the metric
"""
options = None
path = 'projects/%s/metrics/%s' % (project, metric_name)
try:
self._gax_api.delete_log_metric(path, options=options)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(path)
raise
def reload(self):
"""Reload the metadata for this instance.
See:
https://cloud.google.com/spanner/reference/rpc/google.spanner.admin.instance.v1#google.spanner.admin.instance.v1.InstanceAdmin.GetInstanceConfig
"""
api = self._client.instance_admin_api
options = _options_with_prefix(self.name)
try:
instance_pb = api.get_instance(self.name, options=options)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(self.name)
raise
self._update_from_pb(instance_pb)
def logger_delete(self, project, logger_name):
"""API call: delete all entries in a logger via a DELETE request
:type project: str
:param project: ID of project containing the log entries to delete
:type logger_name: str
:param logger_name: name of logger containing the log entries to delete
"""
options = None
path = 'projects/%s/logs/%s' % (project, logger_name)
try:
self._gax_api.delete_log(path, options=options)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(path)
raise
def test_reload_miss_no_project_set(self):
from google.cloud.exceptions import NotFound
access_id = "ACCESS-ID"
connection = mock.Mock(spec=["api_request"])
connection.api_request.side_effect = NotFound("testing")
client = _Client(connection)
metadata = self._make_one(client)
metadata._properties["accessId"] = access_id
with self.assertRaises(NotFound):
metadata.reload()
expected_path = "/projects/{}/hmacKeys/{}".format(
client.DEFAULT_PROJECT, access_id)
expected_kwargs = {"method": "GET", "path": expected_path}
connection.api_request.assert_called_once_with(**expected_kwargs)
def test_bucket_doesnt_exist_can_be_created_on_upload(self, monkeypatch):
task = GCSUpload(bucket="test", create_bucket=True)
client = MagicMock()
client.return_value = MagicMock(get_bucket=MagicMock(
side_effect=NotFound("no bucket")))
monkeypatch.setattr("prefect.tasks.gcp.storage.get_storage_client",
client)
task.run(data="empty", credentials={})
task.run(data="empty", bucket="run", credentials={})
assert client.return_value.create_bucket.called
assert client.return_value.create_bucket.call_args_list[0][0][
0] == "test"
assert client.return_value.create_bucket.call_args_list[1][0][
0] == "run"
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(
request=dict(
parent=TEST_PARENT,
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,
)
self.assertEqual(Instance(name=TEST_NAME), result)
def test_exists_miss(self):
from google.cloud.exceptions import NotFound
client = self._make_client()
bucket = self._make_bucket(client)
notification = self._make_one(bucket, self.TOPIC_NAME)
notification._properties['id'] = self.NOTIFICATION_ID
api_request = client._connection.api_request
api_request.side_effect = NotFound('testing')
self.assertFalse(notification.exists())
api_request.assert_called_once_with(
method='GET',
path=self.NOTIFICATION_PATH,
query_params={},
)
def subscription_delete(self, subscription_path):
"""API call: delete a subscription
See:
https://cloud.google.com/pubsub/docs/reference/rest/v1/projects.subscriptions/delete
:type subscription_path: str
:param subscription_path:
the fully-qualified path of the subscription, in format
``projects/<PROJECT>/subscriptions/<SUB_NAME>``.
"""
try:
self._gax_api.delete_subscription(subscription_path)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(subscription_path)
raise
def subscription_pull(self,
subscription_path,
return_immediately=False,
max_messages=1):
"""API call: retrieve messages for a subscription
See:
https://cloud.google.com/pubsub/docs/reference/rest/v1/projects.subscriptions/modifyPushConfig
:type subscription_path: str
:param subscription_path:
the fully-qualified path of the new subscription, in format
``projects/<PROJECT>/subscriptions/<SUB_NAME>``.
:type return_immediately: bool
:param return_immediately: if True, the back-end returns even if no
messages are available; if False, the API
call blocks until one or more messages are
available.
:type max_messages: int
:param max_messages: the maximum number of messages to return.
:rtype: list of dict
:returns: the ``receivedMessages`` element of the response.
"""
try:
response_pb = self._gax_api.pull(
subscription_path,
max_messages,
return_immediately=return_immediately)
except GaxError as exc:
code = exc_to_code(exc.cause)
if code == StatusCode.NOT_FOUND:
raise NotFound(subscription_path)
elif code == StatusCode.DEADLINE_EXCEEDED:
# NOTE: The JSON-over-HTTP API returns a 200 with an empty
# response when ``return_immediately`` is ``False``, so
# we "mutate" the gRPC error into a non-error to conform.
if not return_immediately:
return []
raise
return [
_received_message_pb_to_mapping(rmpb)
for rmpb in response_pb.received_messages
]
def get_prefixes_list_by_tile_name(self, tile_name):
"""
get prefix list for tile_name from bucket
:param tile_name: str
:return: list
"""
tile_path = self.get_tile_uri(tile_name)
tile_uri = f'L2/tiles/{tile_path}/'
blobs = self.storage_client.list_blobs(self.storage_bucket, prefix=tile_uri, delimiter='/')
blobs._next_page()
prefixes = list(blobs.prefixes)
if not prefixes:
raise NotFound(f'no such tile_uri: {tile_uri}')
return prefixes
def update(self):
"""Update this instance.
See
https://cloud.google.com/spanner/reference/rpc/google.spanner.admin.instance.v1#google.spanner.admin.instance.v1.InstanceAdmin.UpdateInstance
.. note::
Updates the ``display_name`` and ``node_count``. To change those
values before updating, set them via
.. code:: python
instance.display_name = 'New display name'
instance.node_count = 5
before calling :meth:`update`.
:rtype: :class:`google.api_core.operation.Operation`
:returns: an operation instance
:raises NotFound: if the instance does not exist
:raises GaxError: for other errors returned from the call
"""
api = self._client.instance_admin_api
instance_pb = admin_v1_pb2.Instance(
name=self.name,
config=self.configuration_name,
display_name=self.display_name,
node_count=self.node_count,
)
field_mask = FieldMask(paths=['config', 'display_name', 'node_count'])
options = _options_with_prefix(self.name)
try:
future = api.update_instance(
instance=instance_pb,
field_mask=field_mask,
options=options,
)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(self.name)
raise
return future
def delete(self):
"""Delete this session.
See:
https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.Spanner.GetSession
:raises: :exc:`ValueError` if :attr:`session_id` is not already set.
"""
if self._session_id is None:
raise ValueError('Session ID not set by back-end')
api = self._database.spanner_api
options = _options_with_prefix(self._database.name)
try:
api.delete_session(self.name, options=options)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(self.name)
raise
def test_bucket_doesnt_exist_raises_by_default(self, monkeypatch, klass):
task = klass(bucket="test")
run_arg = "data" if isinstance(task, GCSUpload) else "blob"
client = MagicMock(get_bucket=MagicMock(
side_effect=NotFound("no bucket")))
monkeypatch.setattr("prefect.tasks.google.storage.Credentials",
MagicMock())
monkeypatch.setattr(
"prefect.tasks.google.storage.storage.Client",
MagicMock(return_value=client),
)
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(
GOOGLE_APPLICATION_CREDENTIALS={})):
with pytest.raises(NotFound, match="no bucket"):
task.run(**{run_arg: "empty"})
def reload(self):
"""Reload this database.
Refresh any configured schema into :attr:`ddl_statements`.
See
https://cloud.google.com/spanner/reference/rpc/google.spanner.admin.database.v1#google.spanner.admin.database.v1.DatabaseAdmin.GetDatabaseDDL
"""
api = self._instance._client.database_admin_api
options = _options_with_prefix(self.name)
try:
response = api.get_database_ddl(self.name, options=options)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(self.name)
raise
self._ddl_statements = tuple(response.statements)
def topic_list_subscriptions(self, topic, page_size=0, page_token=None):
"""API call: list subscriptions bound to a topic
See:
https://cloud.google.com/pubsub/docs/reference/rest/v1/projects.topics.subscriptions/list
:type topic: :class:`~google.cloud.pubsub.topic.Topic`
:param topic: The topic that owns the subscriptions.
:type page_size: int
:param page_size: maximum number of subscriptions to return, If not
passed, defaults to a value set by the API.
:type page_token: str
:param page_token: opaque marker for the next "page" of subscriptions.
If not passed, the API will return the first page
of subscriptions.
:rtype: :class:`~google.cloud.iterator.Iterator`
:returns: Iterator of
:class:`~google.cloud.pubsub.subscription.Subscription`
accessible to the current API.
:raises: :exc:`~google.cloud.exceptions.NotFound` if the topic does
not exist.
"""
if page_token is None:
page_token = INITIAL_PAGE
options = CallOptions(page_token=page_token)
topic_path = topic.full_name
try:
page_iter = self._gax_api.list_topic_subscriptions(
topic_path, page_size=page_size, options=options)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(topic_path)
raise
iter_kwargs = {}
if page_size: # page_size can be 0 or explicit None.
iter_kwargs['max_results'] = page_size
iterator = GAXIterator(self._client, page_iter, _item_to_subscription,
**iter_kwargs)
iterator.topic = topic
return iterator
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 snapshot_delete(self, snapshot_path):
"""API call: delete a topic
See:
https://cloud.google.com/pubsub/docs/reference/rest/v1/projects.snapshots/delete
:type snapshot_path: str
:param snapshot_path: fully-qualified path of the snapshot, in format
``projects/<PROJECT>/snapshots/<SNAPSHOT_NAME>``.
:raises: :exc:`google.cloud.exceptions.NotFound` if the snapshot does
not exist
"""
try:
self._gax_api.delete_snapshot(snapshot_path)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(snapshot_path)
raise
def test_delete_miss_w_defaults(self):
from google.cloud.exceptions import NotFound
client = mock.Mock(spec=["_delete_resource", "project"])
client._delete_resource.side_effect = NotFound("testing")
client.project = self.BUCKET_PROJECT
bucket = self._make_bucket(client)
notification = self._make_one(bucket, self.TOPIC_NAME)
notification._properties["id"] = self.NOTIFICATION_ID
with self.assertRaises(NotFound):
notification.delete()
client._delete_resource.assert_called_once_with(
self.NOTIFICATION_PATH,
query_params={},
timeout=self._get_default_timeout(),
retry=DEFAULT_RETRY,
)
def test_exists_miss(self):
from google.cloud.exceptions import NotFound
client = self._make_client()
bucket = self._make_bucket(client)
notification = self._make_one(bucket, self.TOPIC_NAME)
notification._properties["id"] = self.NOTIFICATION_ID
api_request = client._connection.api_request
api_request.side_effect = NotFound("testing")
self.assertFalse(notification.exists(timeout=42))
api_request.assert_called_once_with(
method="GET",
path=self.NOTIFICATION_PATH,
query_params={},
timeout=42,
retry=DEFAULT_RETRY,
)
def subscription_acknowledge(self, subscription_path, ack_ids):
"""API call: acknowledge retrieved messages
See:
https://cloud.google.com/pubsub/docs/reference/rest/v1/projects.subscriptions/modifyPushConfig
:type subscription_path: str
:param subscription_path:
the fully-qualified path of the new subscription, in format
``projects/<PROJECT>/subscriptions/<SUB_NAME>``.
:type ack_ids: list of string
:param ack_ids: ack IDs of messages being acknowledged
"""
try:
self._gax_api.acknowledge(subscription_path, ack_ids)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(subscription_path)
raise
def test_validate_error(self):
from google.cloud.exceptions import NotFound
connection = self._make_connection()
# mock snapshot context manager
snapshot_obj = mock.Mock()
snapshot_obj.execute_sql = mock.Mock(side_effect=NotFound("Not found"))
snapshot_ctx = mock.Mock()
snapshot_ctx.__enter__ = mock.Mock(return_value=snapshot_obj)
snapshot_ctx.__exit__ = exit_ctx_func
snapshot_method = mock.Mock(return_value=snapshot_ctx)
connection.database.snapshot = snapshot_method
with self.assertRaises(NotFound):
connection.validate()
snapshot_obj.execute_sql.assert_called_once_with("SELECT 1")
def test_creds_are_pulled_from_secret_at_runtime(self, monkeypatch):
task = CreateBigQueryTable()
creds_loader = MagicMock()
monkeypatch.setattr("prefect.tasks.google.bigquery.Credentials",
creds_loader)
monkeypatch.setattr(
"prefect.tasks.google.bigquery.bigquery.Client",
MagicMock(return_value=MagicMock(get_table=MagicMock(
side_effect=NotFound("boy")))),
)
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(
GOOGLE_APPLICATION_CREDENTIALS={"key": 42})):
task.run()
assert creds_loader.from_service_account_info.call_args[0][0] == {
"key": 42
}
def topic_list_subscriptions(self,
topic_path,
page_size=0,
page_token=None):
"""API call: list subscriptions bound to a topic
See:
https://cloud.google.com/pubsub/docs/reference/rest/v1/projects.topics.subscriptions/list
:type topic_path: str
:param topic_path: fully-qualified path of the topic, in format
``projects/<PROJECT>/topics/<TOPIC_NAME>``.
:type page_size: int
:param page_size: maximum number of subscriptions to return, If not
passed, defaults to a value set by the API.
:type page_token: str
:param page_token: opaque marker for the next "page" of subscriptions.
If not passed, the API will return the first page
of subscriptions.
:rtype: list of strings
:returns: fully-qualified names of subscriptions for the supplied
topic.
:raises: :exc:`google.cloud.exceptions.NotFound` if the topic does not
exist
"""
if page_token is None:
page_token = INITIAL_PAGE
options = CallOptions(page_token=page_token)
try:
page_iter = self._gax_api.list_topic_subscriptions(
topic_path, page_size=page_size, options=options)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(topic_path)
raise
subs = page_iter.next()
token = page_iter.page_token or None
return subs, token
def delete(self):
"""Delete this session.
See
https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.Spanner.GetSession
:raises ValueError: if :attr:`session_id` is not already set.
:raises NotFound: if the session does not exist
:raises GaxError:
for errors other than ``NOT_FOUND`` returned from the call
"""
if self._session_id is None:
raise ValueError('Session ID not set by back-end')
api = self._database.spanner_api
options = _options_with_prefix(self._database.name)
try:
api.delete_session(self.name, options=options)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(self.name)
raise
def subscription_get(self, subscription_path):
"""API call: retrieve a subscription
See:
https://cloud.google.com/pubsub/docs/reference/rest/v1/projects.subscriptions/get
:type subscription_path: str
:param subscription_path:
the fully-qualified path of the subscription, in format
``projects/<PROJECT>/subscriptions/<SUB_NAME>``.
:rtype: dict
:returns: ``Subscription`` resource returned from the API.
"""
try:
sub_pb = self._gax_api.get_subscription(subscription_path)
except GaxError as exc:
if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
raise NotFound(subscription_path)
raise
return _subscription_pb_to_mapping(sub_pb)
def __init__(self, name, mode, bucket):
"""Create a GoogleCloudFile handler.
Args:
name: The name of the file within the bucket (string).
mode: A string of the file mode to open with (string).
bucket: The bucket to load and save to (Bucket).
"""
self.name = name
self.mode = mode
self._mimetype = mimetypes.guess_type(name)[0]
self._is_dirty = False
self._file = None
self._blob = bucket.get_blob(name)
if not self.blob and "w" in mode:
self.blob = Blob(self.name, bucket)
elif not self.blob and "r" in mode:
message = "{} file was not found."
message.format(name)
raise NotFound(message=message)
self.open(mode)
def test_bucket_doesnt_exist_can_be_created_on_upload(self, monkeypatch):
task = GCSUpload(bucket="test", create_bucket=True)
client = MagicMock(get_bucket=MagicMock(
side_effect=NotFound("no bucket")))
monkeypatch.setattr("prefect.tasks.google.storage.Credentials",
MagicMock())
monkeypatch.setattr(
"prefect.tasks.google.storage.storage.Client",
MagicMock(return_value=client),
)
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(
GOOGLE_APPLICATION_CREDENTIALS={})):
task.run(data="empty")
task.run(data="empty", bucket="run")
assert client.create_bucket.called
assert client.create_bucket.call_args_list[0][0][0] == "test"
assert client.create_bucket.call_args_list[1][0][0] == "run"
def test_exists_miss_w_defaults(self):
from google.cloud.exceptions import NotFound
access_id = "ACCESS-ID"
project = "PROJECT"
client = mock.Mock(spec=["_get_resource", "project"])
client._get_resource.side_effect = NotFound("testing")
client.project = project
metadata = self._make_one(client)
metadata._properties["accessId"] = access_id
self.assertFalse(metadata.exists())
expected_path = "/projects/{}/hmacKeys/{}".format(project, access_id)
expected_query_params = {}
client._get_resource.assert_called_once_with(
expected_path,
query_params=expected_query_params,
timeout=self._get_default_timeout(),
retry=DEFAULT_RETRY,
)
def test_update_expire_time_not_found(self):
from google.api_core.exceptions import NotFound
from google.cloud.spanner_admin_database_v1 import Backup
client = _Client()
api = client.database_admin_api = self._make_database_admin_api()
api.update_backup.side_effect = NotFound("testing")
instance = _Instance(self.INSTANCE_NAME, client=client)
backup = self._make_one(self.BACKUP_ID, instance)
expire_time = self._make_timestamp()
with self.assertRaises(NotFound):
backup.update_expire_time(expire_time)
backup_update = Backup(name=self.BACKUP_NAME, expire_time=expire_time,)
update_mask = {"paths": ["expire_time"]}
api.update_backup.assert_called_once_with(
backup=backup_update,
update_mask=update_mask,
metadata=[("google-cloud-resource-prefix", backup.name)],
)
