Skip to content

Commit

Permalink
[FEATURE] TupleGCSStoreBackend::get_all (#9703)
Browse files Browse the repository at this point in the history
  • Loading branch information
tyler-hoffman committed Apr 4, 2024
1 parent 61f8393 commit 5707d2c
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 5 deletions.
20 changes: 15 additions & 5 deletions great_expectations/data_context/store/tuple_store_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -828,12 +828,26 @@ def _build_gcs_object_key(self, key):
return gcs_object_key

def _get(self, key):
gcs_object_key = self._build_gcs_object_key(key)
from great_expectations.compatibility import google

gcs = google.storage.Client(project=self.project)
bucket = gcs.bucket(self.bucket)
return self._get_by_gcs_object_key(bucket, key)

@override
def _get_all(self) -> list[Any]:
from great_expectations.compatibility import google

gcs = google.storage.Client(project=self.project)
bucket = gcs.bucket(self.bucket)

keys = self.list_keys()
keys = [k for k in keys if k != StoreBackend.STORE_BACKEND_ID_KEY]

return [self._get_by_gcs_object_key(bucket, key) for key in keys]

def _get_by_gcs_object_key(self, bucket, key):
gcs_object_key = self._build_gcs_object_key(key)
gcs_response_object = bucket.get_blob(gcs_object_key)
if not gcs_response_object:
raise InvalidKeyError( # noqa: TRY003
Expand All @@ -842,10 +856,6 @@ def _get(self, key):
else:
return gcs_response_object.download_as_bytes().decode("utf-8")

@override
def _get_all(self) -> list[Any]:
raise NotImplementedError

def _set(
self,
key,
Expand Down
64 changes: 64 additions & 0 deletions tests/data_context/store/test_store_backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import pyparsing as pp
import pytest
from moto import mock_s3
from pytest_mock import MockerFixture

from great_expectations.core.data_context_key import DataContextVariableKey
from great_expectations.core.expectation_suite import ExpectationSuite
Expand Down Expand Up @@ -1090,6 +1091,69 @@ def test_TupleGCSStoreBackend(): # noqa: PLR0915
)


@pytest.mark.skipif(
not is_library_loadable(library_name="google.cloud"),
reason="google is not installed",
)
@pytest.mark.skipif(
not is_library_loadable(library_name="google"),
reason="google is not installed",
)
@pytest.mark.big
def test_TupleGCSStoreBackend_get_all(mocker: MockerFixture):
# Note: it would be nice to inject the gcs client so we could pass in the mock.
bucket = "leakybucket"
prefix = "this_is_a_test_prefix"
project = "dummy-project"
base_public_path = "http://www.test.com/"
val_a = "aaa"
val_b = "bbb"

# setup mocks
from great_expectations.compatibility import google

def _create_mock_blob(name: str):
output = mocker.Mock()
output.name = name
return output

def mock_get_blob(gcs_object_key):
"""Test double for bucket::get_blob."""
key_to_return_val = {
f"{prefix}/blob_a": val_a,
f"{prefix}/blob_b": val_b,
}
return mocker.Mock(
download_as_bytes=mocker.Mock(
return_value=key_to_return_val[gcs_object_key].encode("utf-8")
)
)

mock_gcs_client = mocker.MagicMock(spec=google.storage.Client)
mock_gcs_client.list_blobs.return_value = [
_create_mock_blob(name=f"{prefix}/{StoreBackend.STORE_BACKEND_ID_KEY[0]}"),
_create_mock_blob(name=f"{prefix}/blob_a"),
_create_mock_blob(name=f"{prefix}/blob_b"),
]

mock_gcs_client.bucket.return_value = mocker.Mock(
get_blob=mocker.Mock(side_effect=mock_get_blob)
)

with mocker.patch("google.cloud.storage.Client", return_value=mock_gcs_client):
my_store = TupleGCSStoreBackend(
filepath_template=None,
bucket=bucket,
prefix=prefix,
project=project,
base_public_path=base_public_path,
)

result = my_store.get_all()

assert sorted(result) == [val_a, val_b]


@pytest.mark.unit
def test_TupleAzureBlobStoreBackend_credential():
pytest.importorskip("azure.storage.blob")
Expand Down

0 comments on commit 5707d2c

Please sign in to comment.