gen3-dataops-toolkit 2.1.1__tar.gz → 2.1.2__tar.gz

{gen3_dataops_toolkit-2.1.1 → gen3_dataops_toolkit-2.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gen3-dataops-toolkit
-Version: 2.1.1
+Version: 2.1.2
 Summary: Gen3 DataOps toolkit (g3dt): operate SSM-published Gen3 data pipeline environments
 License: Apache-2.0
 Author: JoshuaHarris391

{gen3_dataops_toolkit-2.1.1 → gen3_dataops_toolkit-2.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "gen3-dataops-toolkit"
-version = "2.1.1"
+version = "2.1.2"
 description = "Gen3 DataOps toolkit (g3dt): operate SSM-published Gen3 data pipeline environments"
 authors = ["JoshuaHarris391 <harjo391@gmail.com>"]
 readme = "README.md"

{gen3_dataops_toolkit-2.1.1 → gen3_dataops_toolkit-2.1.2}/src/g3dt/cli/release_cmds.py

@@ -55,6 +55,12 @@ def write(
             f"(workgroup output {rc.athena_output_location})",
             bold=True,
         )
+        # The env's own silver/gold DBs (from SSM) scope the model->DB search:
+        # deterministic in shared accounts, no account-wide Glue perms needed.
+        search_databases = [
+            db for db in (rc.get("glue/db/rawSilver"), rc.get("glue/db/rawGold"))
+            if db
+        ] or None
         release_writer.run(
             dbt_schema_path=dbt_schema_path,
             release_db=rc.release_db,
@@ -65,6 +71,8 @@ def write(
             aws_region=rc.region,
             athena_s3_output=rc.athena_output_location,
             aws_profile=profile,
+            workgroup=rc.athena_workgroup,
+            search_databases=search_databases,
             dry_run=dry_run,
         )
     except config.ConfigError as exc:

{gen3_dataops_toolkit-2.1.1 → gen3_dataops_toolkit-2.1.2}/src/g3dt/utils/athena_utils.py

@@ -26,10 +26,14 @@ class AthenaConfig:
         aws_region (str): AWS region for Athena and S3 operations.
         aws_profile (str): AWS profile to use for authentication.
         athena_s3_output (str): S3 location for Athena query output.
+        workgroup (str, optional): Athena workgroup to run queries in. Roles on
+            the platform are scoped to the env's workgroup, so leaving this
+            unset (awswrangler defaults to 'primary') fails under them.
     """
     aws_region: str
     aws_profile: str
     athena_s3_output: str
+    workgroup: Optional[str] = None
 
     def as_dict(self) -> Dict[str, Any]:
         return {
@@ -87,13 +91,16 @@ class AthenaQuery:
         logger.info(f"Running Athena query: {sql}")
         boto3_session = self._get_boto_session()
         try:
-            df = wr.athena.read_sql_query(
+            kwargs = dict(
                 sql=sql,
                 boto3_session=boto3_session,
                 database=athena_database,
                 ctas_approach=ctas_approach,
-                s3_output=self.config.athena_s3_output
+                s3_output=self.config.athena_s3_output,
             )
+            if getattr(self.config, "workgroup", None):
+                kwargs["workgroup"] = self.config.workgroup
+            df = wr.athena.read_sql_query(**kwargs)
             logger.info(
                 f"Athena query completed successfully. Returned {len(df)} rows."
             )
@@ -159,7 +166,7 @@ class AthenaQuery:
                 database=athena_database,
                 table=table_name,
                 boto3_session=boto3_session,
-                workgroup='primary',
+                workgroup=getattr(self.config, "workgroup", None) or 'primary',
                 temp_path=temp_s3_path
             )
             logger.info(
@@ -172,7 +179,7 @@ class AthenaQuery:
             )
             raise
 
-    def find_db_for_model(self, model_name: str) -> Optional[str]:
+    def find_db_for_model(self, model_name: str, databases: Optional[list] = None) -> Optional[str]:
         """
         Search all Athena databases for a table with the given name and return the database name if found.
 
@@ -202,8 +209,14 @@ class AthenaQuery:
         try:
             # Create a session and pass it to wrangler
             boto3_session = self._get_boto_session()
-            databases = wr.catalog.databases(boto3_session=boto3_session)
-            db_list = databases.get('Database', [])
+            if databases:
+                # Scoped search (e.g. the env's own silver/gold DBs from SSM):
+                # deterministic in shared accounts and needs no account-wide
+                # Glue permissions.
+                db_list = list(databases)
+            else:
+                all_dbs = wr.catalog.databases(boto3_session=boto3_session)
+                db_list = all_dbs.get('Database', [])
             for db in db_list:
                 try:
                     # Pass the session to other wrangler calls too

{gen3_dataops_toolkit-2.1.1 → gen3_dataops_toolkit-2.1.2}/src/g3dt/utils/release_writer.py

@@ -143,6 +143,10 @@ def parse_args():
                         help="S3 URI for Athena query results (e.g., 's3://athena-results-bucket/').")
     parser.add_argument("--release-s3-location", type=str, required=True,
                         help="S3 location for the release table (e.g., 's3://<metadata-bucket>/').")
+    parser.add_argument("--workgroup", type=str, required=False,
+                        help="Athena workgroup to run queries in (platform roles are workgroup-scoped).")
+    parser.add_argument("--search-database", action="append", default=None,
+                        help="Restrict the model->database search to this DB (repeatable).")
     parser.add_argument("--dry-run", action="store_true", default=False,
                         help="Resolve everything and log the SQL; write nothing.")
     parser.add_argument("-v", "--verbose", action="store_true", default=False,
@@ -161,6 +165,8 @@ def run(
     aws_region: str,
     athena_s3_output: str,
     aws_profile: Optional[str] = None,
+    workgroup: Optional[str] = None,
+    search_databases: Optional[list] = None,
     dry_run: bool = False,
 ) -> None:
     """Write one release row per dbt model (idempotent; see insert_release_row).
@@ -176,7 +182,8 @@ def run(
     athena_config = AthenaConfig(
         aws_region=aws_region,
         aws_profile=aws_profile,
-        athena_s3_output=athena_s3_output
+        athena_s3_output=athena_s3_output,
+        workgroup=workgroup,
     )
 
     dbt_models = get_model_names(dbt_schema_path)
@@ -197,7 +204,7 @@ def run(
 
     for model_name in dbt_models:
         logger.info(f"--- Processing model: {model_name}")
-        db_name = athena_query.find_db_for_model(model_name)
+        db_name = athena_query.find_db_for_model(model_name, databases=search_databases)
         if not db_name:
             logger.warning(f"Database not found for model '{model_name}'. Skipping...")
             continue
@@ -239,6 +246,8 @@ def main():
         aws_region=args.aws_region,
         athena_s3_output=args.athena_s3_output,
         aws_profile=args.aws_profile,
+        workgroup=args.workgroup,
+        search_databases=args.search_database,
         dry_run=args.dry_run,
     )