Skip to content

Commit 1b555ce

Browse files
NathanFarmerbilldirks
NathanFarmer
and
billdirks
authored
[MAINTENANCE] Add workspace_id to store_backend dict (great-expectations#11371)
Co-authored-by: Bill Dirks <bill@greatexpectations.io>
1 parent 962eb5a commit 1b555ce

22 files changed

+259
-102
lines changed

‎great_expectations/data_context/data_context/cloud_data_context.py‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@
6767
from great_expectations.alias_types import PathStr
6868
from great_expectations.checkpoint.checkpoint import Checkpoint, CheckpointResult
6969
from great_expectations.core.suite_parameters import SuiteParameterDict
70+
from great_expectations.data_context.store.store import StoreConfigTypedDict
7071
from great_expectations.datasource.fluent import Datasource as FluentDatasource
7172
from great_expectations.render.renderer.site_builder import SiteBuilder
7273

@@ -335,7 +336,7 @@ def retrieve_data_context_config_from_cloud(
335336
:return: the configuration object retrieved from the Cloud API
336337
""" # noqa: E501 # FIXME CoP
337338
response = cls._request_cloud_backend(
338-
cloud_config=cloud_config, resource="data_context_configuration"
339+
cloud_config=cloud_config, resource=GXCloudRESTResource.DATA_CONTEXT
339340
)
340341
config = cls._prepare_v1_config(config=response.json())
341342
return DataContextConfig(**config)
@@ -429,7 +430,7 @@ def _request_cloud_backend(cls, cloud_config: GXCloudConfig, resource: str) -> R
429430
raise OrganizationIdNotSpecifiedError()
430431

431432
with create_session(access_token=access_token) as session:
432-
if resource in [GXCloudRESTResource.ACCOUNTS_ME, GXCloudRESTResource.DATA_CONTEXT]:
433+
if resource == GXCloudRESTResource.ACCOUNTS_ME:
433434
url_workspace_id = None
434435
else:
435436
url_workspace_id = workspace_id
@@ -562,6 +563,14 @@ def _init_datasources(self) -> None:
562563
datasource_store=self._datasource_store,
563564
)
564565

566+
@override
567+
def _init_primary_stores(self, store_configs: Dict[str, StoreConfigTypedDict]) -> None:
568+
for store_config in store_configs.values():
569+
store_config.get("store_backend", {}).get("ge_cloud_credentials", {}).setdefault(
570+
"workspace_id", self.ge_cloud_config.workspace_id
571+
)
572+
super()._init_primary_stores(store_configs)
573+
565574
@override
566575
def _init_datasource_store(self) -> DatasourceStore:
567576
# Never explicitly referenced but adheres

‎great_expectations/data_context/data_context_variables.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -386,6 +386,7 @@ def _init_store(self) -> DataContextStore:
386386
"ge_cloud_credentials": {
387387
"access_token": self.ge_cloud_access_token,
388388
"organization_id": self.ge_cloud_organization_id,
389+
"workspace_id": self.ge_cloud_workspace_id,
389390
},
390391
"suppress_store_backend_id": True,
391392
}

‎tests/analytics/test_analytics.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,11 @@ def test_ephemeral_context_init_with_optional_fields(monkeypatch):
147147
@pytest.mark.parametrize("user_agent_str", [None, "test / x.x.x"])
148148
@pytest.mark.cloud
149149
def test_cloud_context_init(
150-
user_agent_str: Optional[str], cloud_api_fake, cloud_details, monkeypatch
150+
unset_gx_env_variables: None,
151+
user_agent_str: Optional[str],
152+
cloud_api_fake,
153+
cloud_details,
154+
monkeypatch,
151155
):
152156
monkeypatch.setattr(ENV_CONFIG, "gx_analytics_enabled", True) # Enable usage stats
153157

‎tests/core/factory/test_checkpoint_factory.py‎

Lines changed: 42 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,10 @@ def test_checkpoint_factory_add_success_filesystem(empty_data_context):
165165

166166

167167
@pytest.mark.cloud
168-
def test_checkpoint_factory_add_success_cloud(empty_cloud_context_fluent):
168+
def test_checkpoint_factory_add_success_cloud(
169+
unset_gx_env_variables: None,
170+
empty_cloud_context_fluent,
171+
):
169172
_test_checkpoint_factory_add_success(empty_cloud_context_fluent)
170173

171174

@@ -201,7 +204,10 @@ def test_checkpoint_factory_delete_success_filesystem(empty_data_context):
201204

202205

203206
@pytest.mark.cloud
204-
def test_checkpoint_factory_delete_success_cloud(empty_cloud_context_fluent):
207+
def test_checkpoint_factory_delete_success_cloud(
208+
unset_gx_env_variables: None,
209+
empty_cloud_context_fluent,
210+
):
205211
_test_checkpoint_factory_delete_success(empty_cloud_context_fluent)
206212

207213

@@ -346,7 +352,11 @@ def test_checkpoint_factory_add_emits_event_filesystem(self, empty_data_context)
346352
self._test_checkpoint_factory_add_emits_event(empty_data_context)
347353

348354
@pytest.mark.cloud
349-
def test_checkpoint_factory_add_emits_event_cloud(self, empty_cloud_context_fluent):
355+
def test_checkpoint_factory_add_emits_event_cloud(
356+
self,
357+
unset_gx_env_variables: None,
358+
empty_cloud_context_fluent,
359+
):
350360
self._test_checkpoint_factory_add_emits_event(empty_cloud_context_fluent)
351361

352362
def _test_checkpoint_factory_add_emits_event(self, context):
@@ -385,7 +395,11 @@ def test_checkpoint_factory_delete_emits_event_filesystem(self, empty_data_conte
385395
self._test_checkpoint_factory_delete_emits_event(empty_data_context)
386396

387397
@pytest.mark.cloud
388-
def test_checkpoint_factory_delete_emits_event_cloud(self, empty_cloud_context_fluent):
398+
def test_checkpoint_factory_delete_emits_event_cloud(
399+
self,
400+
unset_gx_env_variables: None,
401+
empty_cloud_context_fluent,
402+
):
389403
self._test_checkpoint_factory_delete_emits_event(empty_cloud_context_fluent)
390404

391405
def _test_checkpoint_factory_delete_emits_event(self, context):
@@ -421,7 +435,11 @@ def _test_checkpoint_factory_delete_emits_event(self, context):
421435
class TestCheckpointFactoryAddOrUpdate:
422436
CHECKPOINT_NAME = "checkpoint A"
423437

424-
def test_add_empty_new_checkpoint(self, data_context: AbstractDataContext) -> None:
438+
def test_add_empty_new_checkpoint(
439+
self,
440+
unset_gx_env_variables: None,
441+
data_context: AbstractDataContext,
442+
) -> None:
425443
# arrange
426444
checkpoint = Checkpoint(name=self.CHECKPOINT_NAME, validation_definitions=[])
427445

@@ -432,7 +450,11 @@ def test_add_empty_new_checkpoint(self, data_context: AbstractDataContext) -> No
432450
assert created_checkpoint.id
433451
data_context.checkpoints.get(self.CHECKPOINT_NAME)
434452

435-
def test_add_new_checkpoint_with_validations(self, data_context: AbstractDataContext) -> None:
453+
def test_add_new_checkpoint_with_validations(
454+
self,
455+
unset_gx_env_variables: None,
456+
data_context: AbstractDataContext,
457+
) -> None:
436458
# arrange
437459
batch_def = (
438460
data_context.data_sources.add_pandas("data source A")
@@ -471,7 +493,9 @@ def test_add_new_checkpoint_with_validations(self, data_context: AbstractDataCon
471493
assert val_def_dict == created_val_def.dict()
472494

473495
def test_update_existing_checkpoint_adds_validations(
474-
self, data_context: AbstractDataContext
496+
self,
497+
unset_gx_env_variables: None,
498+
data_context: AbstractDataContext,
475499
) -> None:
476500
# arrange
477501
data_context.checkpoints.add(
@@ -514,7 +538,9 @@ def test_update_existing_checkpoint_adds_validations(
514538
assert val_def_dict == created_val_def.dict()
515539

516540
def test_update_existing_checkpoint_updates_validations(
517-
self, data_context: AbstractDataContext
541+
self,
542+
unset_gx_env_variables: None,
543+
data_context: AbstractDataContext,
518544
) -> None:
519545
# arrange
520546
batch_def = (
@@ -556,7 +582,9 @@ def test_update_existing_checkpoint_updates_validations(
556582
assert val_def.suite.name == new_suite_name
557583

558584
def test_update_existing_checkpoint_deletes_validations(
559-
self, data_context: AbstractDataContext
585+
self,
586+
unset_gx_env_variables: None,
587+
data_context: AbstractDataContext,
560588
) -> None:
561589
# arrange
562590
batch_def = (
@@ -609,7 +637,11 @@ def test_update_existing_checkpoint_deletes_validations(
609637
assert created_checkpoint.validation_definitions[0].name == VALIDATION_DEFINITION_NAME
610638
assert created_checkpoint.validation_definitions[0].suite.name == SUITE_NAME
611639

612-
def test_add_or_update_is_idempotent(self, data_context: AbstractDataContext) -> None:
640+
def test_add_or_update_is_idempotent(
641+
self,
642+
unset_gx_env_variables: None,
643+
data_context: AbstractDataContext,
644+
) -> None:
613645
# arrange
614646
batch_def = (
615647
data_context.data_sources.add_pandas("data source A")

‎tests/core/factory/test_suite_factory.py‎

Lines changed: 40 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,10 @@ def test_suite_factory_delete_success_filesystem(empty_data_context):
192192

193193

194194
@pytest.mark.cloud
195-
def test_suite_factory_delete_success_cloud(empty_cloud_context_fluent):
195+
def test_suite_factory_delete_success_cloud(
196+
unset_gx_env_variables: None,
197+
empty_cloud_context_fluent,
198+
):
196199
_test_suite_factory_delete_success(empty_cloud_context_fluent)
197200

198201

@@ -220,7 +223,11 @@ def _test_suite_factory_delete_success(context):
220223
pytest.param("empty_data_context", id="filesystem", marks=pytest.mark.filesystem),
221224
],
222225
)
223-
def test_suite_factory_all(context_fixture_name: str, request: pytest.FixtureRequest):
226+
def test_suite_factory_all(
227+
unset_gx_env_variables: None,
228+
context_fixture_name: str,
229+
request: pytest.FixtureRequest,
230+
):
224231
context: AbstractDataContext = request.getfixturevalue(context_fixture_name)
225232

226233
# Arrange
@@ -334,7 +341,11 @@ def test_suite_factory_all_with_bad_pydantic_config(
334341

335342

336343
class TestSuiteFactoryAddOrUpdate:
337-
def test_add_empty_new_suite(self, data_context: AbstractDataContext) -> None:
344+
def test_add_empty_new_suite(
345+
self,
346+
unset_gx_env_variables: None,
347+
data_context: AbstractDataContext,
348+
) -> None:
338349
# arrange
339350
suite_name = "suite A"
340351
suite = ExpectationSuite(name=suite_name)
@@ -346,7 +357,11 @@ def test_add_empty_new_suite(self, data_context: AbstractDataContext) -> None:
346357
assert created_suite.id
347358
data_context.suites.get(suite_name)
348359

349-
def test_add_new_suite_with_expectations(self, data_context: AbstractDataContext) -> None:
360+
def test_add_new_suite_with_expectations(
361+
self,
362+
unset_gx_env_variables: None,
363+
data_context: AbstractDataContext,
364+
) -> None:
350365
# arrange
351366
suite_name = "suite A"
352367
expectations = [
@@ -377,7 +392,9 @@ def test_add_new_suite_with_expectations(self, data_context: AbstractDataContext
377392
assert exp == created_exp
378393

379394
def test_update_existing_suite_adds_expectations(
380-
self, data_context: AbstractDataContext
395+
self,
396+
unset_gx_env_variables: None,
397+
data_context: AbstractDataContext,
381398
) -> None:
382399
# arrange
383400
suite_name = "suite A"
@@ -454,7 +471,9 @@ def _test_update_existing_suite_updates_expectations(
454471
assert old_exp.id != new_exp.id
455472

456473
def test_update_existing_suite_deletes_expectations(
457-
self, data_context: AbstractDataContext
474+
self,
475+
unset_gx_env_variables: None,
476+
data_context: AbstractDataContext,
458477
) -> None:
459478
# arrange
460479
suite_name = "suite A"
@@ -490,7 +509,11 @@ def test_update_existing_suite_deletes_expectations(
490509
assert updated_suite.id == existing_suite.id
491510
assert updated_suite.expectations == []
492511

493-
def test_add_or_update_is_idempotent(self, data_context: AbstractDataContext) -> None:
512+
def test_add_or_update_is_idempotent(
513+
self,
514+
unset_gx_env_variables: None,
515+
data_context: AbstractDataContext,
516+
) -> None:
494517
# arrange
495518
suite_name = "suite A"
496519
expectations = [
@@ -519,7 +542,11 @@ def test_add_or_update_is_idempotent(self, data_context: AbstractDataContext) ->
519542

520543

521544
class TestSuiteFactoryAnalytics:
522-
def test_suite_factory_add_emits_event(self, data_context: AbstractDataContext) -> None:
545+
def test_suite_factory_add_emits_event(
546+
self,
547+
unset_gx_env_variables: None,
548+
data_context: AbstractDataContext,
549+
) -> None:
523550
# Arrange
524551
name = "test-suite"
525552
suite = ExpectationSuite(name=name)
@@ -535,7 +562,11 @@ def test_suite_factory_add_emits_event(self, data_context: AbstractDataContext)
535562
event=ExpectationSuiteCreatedEvent(expectation_suite_id=suite.id)
536563
)
537564

538-
def test_suite_factory_delete_emits_event(self, data_context: AbstractDataContext) -> None:
565+
def test_suite_factory_delete_emits_event(
566+
self,
567+
unset_gx_env_variables: None,
568+
data_context: AbstractDataContext,
569+
) -> None:
539570
# Arrange
540571
name = "test-suite"
541572
suite = ExpectationSuite(name=name)

0 commit comments

Comments
 (0)