Merge pull request #170 from broadinstitute/sn_add_more_permission_ch…

…ecks

Check permissions for specific dataset

python/utils/bq_utils.py

@@ -57,7 +57,32 @@ def query_table(self, query: str, to_dataframe: bool = False) -> Any:
             return query_job.result().to_dataframe()
         return [row for row in query_job.result()]
 
-    def check_permissions(self, raise_on_other_failure: bool = True) -> bool:
+    def check_permissions_to_project(self, raise_on_other_failure: bool = True) -> bool:
+        """
+        Checks if the user has permission to access the project.
+
+        Args:
+            raise_on_other_failure (bool): If True, raises an error if an unexpected error occurs. Default is True.
+
+        Returns:
+            bool: True if the user has permissions, False if a 403 Forbidden error is encountered.
+        """
+        return self._check_permissions("SELECT 1", raise_on_other_failure)
+
+    def check_permissions_for_query(self, query: str, raise_on_other_failure: bool = True) -> bool:
+        """
+        Checks if the user has permission to run a specific query.
+
+        Args:
+            query (str): SQL query to execute.
+            raise_on_other_failure (bool): If True, raises an error if an unexpected error occurs. Default is True.
+
+        Returns:
+            bool: True if the user has permissions, False if a 403 Forbidden error is encountered.
+        """
+        return self._check_permissions(query, raise_on_other_failure)
+
+    def _check_permissions(self, qry: str, raise_on_other_failure: bool = True) -> bool:
         """
         Checks if the user has permission to run queries and access the project.
 
@@ -69,7 +94,7 @@ def check_permissions(self, raise_on_other_failure: bool = True) -> bool:
         """
         try:
             # A simple query that should succeed if the user has permissions
-            query = "SELECT 1"
+            query = qry
             self.client.query(query).result()  # Run a lightweight query
             return True
         except Forbidden:

python/utils/tdr_utils/tdr_bq_utils.py

@@ -6,7 +6,7 @@
 
 
 class GetTdrAssetInfo:
-    def __init__(self, tdr: TDR, dataset_id: Optional[str], snapshot_id: Optional[str]):
+    def __init__(self, tdr: TDR, dataset_id: Optional[str] = None, snapshot_id: Optional[str] = None):
         """
         Initialize the GetTdrAssetInfo class.
 
@@ -15,6 +15,8 @@ def __init__(self, tdr: TDR, dataset_id: Optional[str], snapshot_id: Optional[st
             dataset_id (Optional[str]): ID of the dataset.
             snapshot_id (Optional[str]): ID of the snapshot.
         """
+        if not dataset_id and not snapshot_id:
+            raise ValueError("Either dataset_id or snapshot_id must be provided.")
         self.tdr = tdr
         self.dataset_id = dataset_id
         self.snapshot_id = snapshot_id
@@ -80,17 +82,21 @@ def __init__(self, project_id: str, bq_schema: str):
         self.bq_schema = bq_schema
         self.bq_util = BigQueryUtil(project_id)
 
-    def check_permissions(self, raise_on_other_failure: bool) -> bool:
+    def check_permissions_for_dataset(self, raise_on_other_failure: bool) -> bool:
         """
-        Check the permissions for accessing BigQuery.
+        Check the permissions for accessing BigQuery for specific dataset.
 
         Args:
             raise_on_other_failure (bool): Whether to raise an exception on other failures.
 
         Returns:
             bool: True if permissions are sufficient, False otherwise.
         """
-        return self.bq_util.check_permissions(raise_on_other_failure=raise_on_other_failure)
+        query = f"""SELECT 1 FROM `{self.project_id}.{self.bq_schema}.INFORMATION_SCHEMA.TABLES`"""
+        return self.bq_util.check_permissions_for_query(
+            query=query,
+            raise_on_other_failure=raise_on_other_failure
+        )
 
     def get_tdr_table_contents(self, exclude_datarepo_id: bool, table_name: str, to_dataframe: bool) -> Any:
         """