import logging
import os
from google.cloud import bigquery
from google.cloud.bigquery.dataset import AccessEntry
from dotenv import load_dotenv

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)

# Load environment variables from a .env file
load_dotenv()

# Retrieve environment variables
SERVICE_ACCOUNT_JSON = os.getenv("GCP_C2C_SERVICE_ACCOUNT_JSON")
PROJECT_ID = os.getenv("GCP_C2C_PROJECT_ID")
DATASET_ID = "jsons_rep_data_warehouse"
MASKED_READER_EMAIL = os.getenv("MASKEDREADER_EMAIL")
FINEGRAINEDREADER_EMAIL = os.getenv("FINEGRAINEDREADER_EMAIL")


# TODO(developer): Set dataset_id to the ID of the dataset to fetch.
dataset_id = f'{PROJECT_ID}.{DATASET_ID}'

# TODO(developer): Set entity_id to the ID of the email or group from whom
# you are adding access. Alternatively, to the JSON REST API representation
# of the entity, such as a view's table reference.
entity_id = MASKED_READER_EMAIL

from google.cloud.bigquery.enums import EntityTypes

# TODO(developer): Set entity_type to the type of entity you are granting access to.
# Common types include:
#
# * "userByEmail" -- A single user or service account. For example "[email protected]"
# * "groupByEmail" -- A group of users. For example "[email protected]"
# * "view" -- An authorized view. For example
#       {"projectId": "p", "datasetId": "d", "tableId": "v"}
#
# For a complete reference, see the REST API reference documentation:
# https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#Dataset.FIELDS.access
entity_type_group = EntityTypes.GROUP_BY_EMAIL
entity_type_user = EntityTypes.USER_BY_EMAIL
entity_type_view = EntityTypes.VIEW

# TODO(developer): Set role to a one of the "Basic roles for datasets"
# described here:
# https://cloud.google.com/bigquery/docs/access-control-basic-roles#dataset-basic-roles
role = "READER"

from google.cloud import bigquery

# Construct a BigQuery client object.
client = bigquery.Client.from_service_account_json(SERVICE_ACCOUNT_JSON)

dataset = client.get_dataset(dataset_id)  # Make an API request.

entries = list(dataset.access_entries)
entries.append(
    bigquery.AccessEntry(
        role=role,
        entity_type=entity_type_user,
        entity_id=entity_id,
    )
)
dataset.access_entries = entries

dataset = client.update_dataset(dataset, ["access_entries"])  # Make an API request.

full_dataset_id = "{}.{}".format(dataset.project, dataset.dataset_id)
print(
    "Updated dataset '{}' with modified user permissions.".format(full_dataset_id)
)

(bigquery) nyck33@lenovo-gtx1650:/mnt/d/c2c/sdp-masking-nov28-2024/policy_tags_udfs$ python set_iam_update_dataset_access_control_sample.py
Traceback (most recent call last):
  File "/mnt/d/c2c/sdp-masking-nov28-2024/policy_tags_udfs/set_iam_update_dataset_access_control_sample.py", line 71, in <module>
    dataset = client.update_dataset(dataset, ["access_entries"])  # Make an API request.
  File "/home/nyck33/miniconda3/envs/bigquery/lib/python3.9/site-packages/google/cloud/bigquery/client.py", line 1228, in update_dataset
    api_response = self._call_api(
  File "/home/nyck33/miniconda3/envs/bigquery/lib/python3.9/site-packages/google/cloud/bigquery/client.py", line 833, in _call_api
    return call()
  File "/home/nyck33/miniconda3/envs/bigquery/lib/python3.9/site-packages/google/api_core/retry/retry_unary.py", line 293, in retry_wrapped_func
    return retry_target(
  File "/home/nyck33/miniconda3/envs/bigquery/lib/python3.9/site-packages/google/api_core/retry/retry_unary.py", line 153, in retry_target
    _retry_error_helper(
  File "/home/nyck33/miniconda3/envs/bigquery/lib/python3.9/site-packages/google/api_core/retry/retry_base.py", line 212, in _retry_error_helper
    raise final_exc from source_exc
  File "/home/nyck33/miniconda3/envs/bigquery/lib/python3.9/site-packages/google/api_core/retry/retry_unary.py", line 144, in retry_target
    result = target()
  File "/home/nyck33/miniconda3/envs/bigquery/lib/python3.9/site-packages/google/cloud/_http/__init__.py", line 494, in api_request
    raise exceptions.from_http_response(response)
google.api_core.exceptions.BadRequest: 400 PATCH https://bigquery.googleapis.com/bigquery/v2/projects/c2c-dwh-dev/datasets/jsons_rep_data_warehouse?prettyPrint=false: An access entry must have exactly one of userByEmail, groupByEmail, domain, specialGroup defined, view, routine, or dataset.
(bigquery) nyck33@lenovo-gtx1650:/mnt/d/c2c/sdp-masking-nov28-2024/policy_tags_udfs$