Skip to content

Commit

Permalink
implement funcitonal test
Browse files Browse the repository at this point in the history
  • Loading branch information
salimmoulouel committed May 3, 2024
1 parent 808d41b commit 18e55ce
Show file tree
Hide file tree
Showing 5 changed files with 62 additions and 10 deletions.
2 changes: 1 addition & 1 deletion dbt/adapters/bigquery/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ def add_access_entry_to_dataset(dataset: Dataset, access_entry: AccessEntry) ->



def delete_access_entry_to_dataset(dataset: Dataset, access_entry: AccessEntry) -> Dataset:
def delete_access_entry_from_dataset(dataset: Dataset, access_entry: AccessEntry) -> Dataset:
"""Remove an access entry from a dataset, always use.
Args:
Expand Down
12 changes: 6 additions & 6 deletions dbt/adapters/bigquery/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@
from dbt.adapters.bigquery import BigQueryColumn, BigQueryConnectionManager
from dbt.adapters.bigquery.column import get_nested_column_data_types
from dbt.adapters.bigquery.connections import BigQueryAdapterResponse
from dbt.adapters.bigquery.dataset import add_access_entry_to_dataset, is_access_entry_in_dataset, delete_access_entry_to_dataset
from dbt.adapters.bigquery.dataset import add_access_entry_to_dataset, is_access_entry_in_dataset, delete_access_entry_from_dataset
from dbt.adapters.bigquery.python_submissions import (
ClusterDataprocHelper,
ServerlessDataProcHelper,
Expand Down Expand Up @@ -840,17 +840,17 @@ def grant_access_to(self, entity, entity_type, role, grant_target_dict,full_refr
logger.warning(f"Access entry {access_entry} " f"already exists in dataset")
return
else:
dataset = delete_access_entry_to_dataset(dataset,access_entry)
dataset = delete_access_entry_from_dataset(dataset,access_entry)
dataset = client.update_dataset(
dataset,
["access_entries"],
) # Make an API request.
full_dataset_id = f"{dataset.project}.{dataset.dataset_id}"
print(f"Revoked dataset access for '{access_entry.entity_id}' to ' dataset '{full_dataset_id}.'")
logger.info(f"Revoked dataset access for '{access_entry.entity_id}' to ' dataset '{full_dataset_id}.'")
dataset = add_access_entry_to_dataset(dataset, access_entry)
dataset = client.update_dataset(dataset, ["access_entries"])
full_dataset_id = f"{dataset.project}.{dataset.dataset_id}"
print(f"allowed dataset access for '{access_entry.entity_id}' to ' dataset '{full_dataset_id}.'")
logger.info(f"allowed dataset access for '{access_entry.entity_id}' to ' dataset '{full_dataset_id}.'")


@available.parse_none
Expand All @@ -870,14 +870,14 @@ def remove_grant_access_to(self, entity, entity_type, role, grant_target_dict):
access_entry = AccessEntry(role, entity_type, entity)
# only perform removing if access entry in dataset
if is_access_entry_in_dataset(dataset, access_entry):
dataset = delete_access_entry_to_dataset(dataset,access_entry)
dataset = delete_access_entry_from_dataset(dataset,access_entry)
dataset = client.update_dataset(
dataset,
["access_entries"],
) # Make an API request.

full_dataset_id = f"{dataset.project}.{dataset.dataset_id}"
print(f"Revoked dataset access for '{access_entry.entity_id}' to ' dataset '{full_dataset_id}.'")
logger.info(f"Revoked dataset access for '{access_entry.entity_id}' to ' dataset '{full_dataset_id}.'")
else:
logger.warning(f"Access entry {access_entry} not in the dataset {full_dataset_id} no need to remove it")

Expand Down
2 changes: 1 addition & 1 deletion dbt/include/bigquery/macros/materializations/view.sql
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@

{% if config.get('grant_access_to') %}
{% for grant_target_dict in config.get('grant_access_to') %}
{% do adapter.grant_access_to(this, 'view', None, grant_target_dict,should_full_refresh()) %}
{% do adapter.grant_access_to(this, 'view', None, grant_target_dict, should_full_refresh()) %}
{% endfor %}
{% endif %}

Expand Down
52 changes: 52 additions & 0 deletions tests/functional/adapter/test_grant_access_to.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,58 @@ def test_grant_access_succeeds(self, project, setup_grant_schema, teardown_grant
assert len(results) == 2



class TestAccessGrantSucceedsWithFullRefresh:
@pytest.fixture(scope="class")
def setup_grant_schema(
self,
project,
unique_schema,
):
with project.adapter.connection_named("__test_grants"):
relation = project.adapter.Relation.create(
database=project.database,
schema=get_schema_name(unique_schema),
identifier="grant_access",
)
project.adapter.create_schema(relation)
yield relation

@pytest.fixture(scope="class")
def teardown_grant_schema(
self,
project,
unique_schema,
):
yield
with project.adapter.connection_named("__test_grants"):
relation = project.adapter.Relation.create(
database=project.database, schema=get_schema_name(unique_schema)
)
project.adapter.drop_schema(relation)

@pytest.fixture(scope="class")
def models(self, unique_schema):
dataset = get_schema_name(unique_schema)
return {
"select_1.sql": select_1(dataset=dataset, materialized="view"),
"select_1_table.sql": select_1(dataset=dataset, materialized="table"),
}

def test_grant_access_succeeds(self, project, setup_grant_schema, teardown_grant_schema):
# Need to run twice to validate idempotency
results = run_dbt(["run"])
assert len(results) == 2
time.sleep(10)
results = run_dbt(["run","--full-refresh"])
assert len(results) == 2
time.sleep(10)
results = run_dbt(["run"])
assert len(results) == 2




class TestAccessGrantFails:
@pytest.fixture(scope="class")
def models(self):
Expand Down
4 changes: 2 additions & 2 deletions tests/unit/test_dataset.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from dbt.adapters.bigquery.dataset import add_access_entry_to_dataset, is_access_entry_in_dataset, delete_access_entry_to_dataset
from dbt.adapters.bigquery.dataset import add_access_entry_to_dataset, is_access_entry_in_dataset, delete_access_entry_from_dataset
from dbt.adapters.bigquery import BigQueryRelation

from google.cloud.bigquery import Dataset, AccessEntry, DatasetReference
Expand Down Expand Up @@ -114,5 +114,5 @@ def test_delete_access_to_dataset_updates_dataset():
access_entry = AccessEntry(None, "table", entity)
dataset = add_access_entry_to_dataset(dataset, access_entry)
assert is_access_entry_in_dataset(dataset, access_entry)
dataset = delete_access_entry_to_dataset(dataset, access_entry)
dataset = delete_access_entry_from_dataset(dataset, access_entry)
assert not is_access_entry_in_dataset(dataset, access_entry)

0 comments on commit 18e55ce

Please sign in to comment.