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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gen3-dataops-toolkit
-Version: 2.1.0
+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.0 → gen3_dataops_toolkit-2.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "gen3-dataops-toolkit"
-version = "2.1.0"
+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.0 → 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.0 → gen3_dataops_toolkit-2.1.2}/src/g3dt/config.py

@@ -336,7 +336,9 @@ def _studies_from_s3_cached(project: str, base_env: str, profile: Optional[str])
     from g3dt import resolver
 
     rc = resolver.resolve(project, base_env, profile=profile)
-    session = boto3.Session(profile_name=profile) if profile else boto3.Session()
+    # Explicit region: botocore ignores AWS_REGION (only AWS_DEFAULT_REGION),
+    # so an ambient session on the EC2 box would have no region at all.
+    session = boto3.Session(profile_name=profile, region_name=rc.region)
     try:
         body = session.client("s3").get_object(
             Bucket=rc.metadata_bucket, Key=STUDIES_S3_KEY

{gen3_dataops_toolkit-2.1.0 → gen3_dataops_toolkit-2.1.2}/src/g3dt/resolver.py

@@ -127,31 +127,55 @@ class ResolvedConfig:
         return self.params.get(key, default)
 
 
+def _session(profile: Optional[str], region: Optional[str]) -> boto3.Session:
+    """A boto3 session with an EXPLICIT region.
+
+    The region always comes from the marker (file value, or its AWS_REGION
+    env-var override) rather than boto3's own chain: botocore only honours
+    AWS_DEFAULT_REGION, so on the EC2 job box — where user-data exports
+    AWS_REGION and there is no ~/.aws/config — an ambient-chain session has
+    no region at all (NoRegionError).
+    """
+    if region is None:
+        from g3dt.config import load_marker  # late import (config imports us)
+
+        region = load_marker()["region"]
+    return boto3.Session(profile_name=profile, region_name=region)
+
+
 @functools.lru_cache(maxsize=None)
-def resolve(project: str, env: str, profile: Optional[str] = None) -> ResolvedConfig:
+def resolve(
+    project: str,
+    env: str,
+    profile: Optional[str] = None,
+    region: Optional[str] = None,
+) -> ResolvedConfig:
     """Resolve all deployed names for ``project``/``env`` from SSM (cached).
 
-    Cached on ``(project, env, profile)`` for the life of the process, so a
-    single CLI invocation makes exactly one ``get_parameters_by_path``
+    Cached on ``(project, env, profile, region)`` for the life of the process,
+    so a single CLI invocation makes exactly one ``get_parameters_by_path``
     round-trip no matter how many commands ask for a name. Call
     ``resolve.cache_clear()`` in tests.
 
     ``profile`` is the local AWS named profile to authenticate the read (e.g.
     ``etl_test``); ``None`` uses the default credential chain — which is the
     instance profile on the EC2 job box and the build role in CodeBuild.
+    ``region`` defaults to the marker's region.
     """
-    session = boto3.Session(profile_name=profile) if profile else boto3.Session()
+    session = _session(profile, region)
     return ResolvedConfig(
         project=project, env=env, params=_fetch_params(project, env, session=session)
     )
 
 
-def list_envs(project: str, profile: Optional[str] = None) -> list:
+def list_envs(
+    project: str, profile: Optional[str] = None, region: Optional[str] = None
+) -> list:
     """Return the environment names that have an SSM tree under /{project}/.
 
     Reads the parameter names one level down (e.g. ``/etl/test/...`` -> ``test``).
     """
-    session = boto3.Session(profile_name=profile) if profile else boto3.Session()
+    session = _session(profile, region)
     client = session.client("ssm")
     paginator = client.get_paginator("get_parameters_by_path")
     envs = set()

{gen3_dataops_toolkit-2.1.0 → 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.0 → 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,
     )