Add workgroup global config. #437

Author: igorborgest

awswrangler/_config.py

@@ -30,6 +30,7 @@ class _ConfigArg(NamedTuple):
     "max_cache_query_inspections": _ConfigArg(dtype=int, nullable=False),
     "max_cache_seconds": _ConfigArg(dtype=int, nullable=False),
     "s3_block_size": _ConfigArg(dtype=int, nullable=False, enforced=True),
+    "workgroup": _ConfigArg(dtype=str, nullable=False, enforced=True),
 }
 
 
@@ -214,6 +215,15 @@ def s3_block_size(self) -> int:
     def s3_block_size(self, value: int) -> None:
         self._set_config_value(key="s3_block_size", value=value)
 
+    @property
+    def workgroup(self) -> Optional[str]:
+        """Property workgroup."""
+        return cast(Optional[str], self["workgroup"])
+
+    @workgroup.setter
+    def workgroup(self, value: Optional[str]) -> None:
+        self._set_config_value(key="workgroup", value=value)
+
 
 def _inject_config_doc(doc: Optional[str], available_configs: Tuple[str, ...]) -> str:
     if doc is None:

tests/test_config.py

@@ -9,7 +9,7 @@
 logging.getLogger("awswrangler").setLevel(logging.DEBUG)
 
 
-def test_basics(path, glue_database, glue_table):
+def test_basics(path, glue_database, glue_table, workgroup0, workgroup1):
     args = {"table": glue_table, "path": "", "columns_types": {"col0": "bigint"}}
 
     # Missing database argument
@@ -62,3 +62,12 @@ def test_basics(path, glue_database, glue_table):
         wr.catalog.does_table_exist(table=glue_table)
 
     assert wr.config.to_pandas().shape == (len(wr._config._CONFIG_ARGS), 7)
+
+    # Workgroup
+    wr.config.workgroup = workgroup0
+    df = wr.athena.read_sql_query(sql="SELECT 1 as col0", database=glue_database)
+    assert df.query_metadata["WorkGroup"] == workgroup0
+    os.environ["WR_WORKGROUP"] = workgroup1
+    wr.config.reset()
+    df = wr.athena.read_sql_query(sql="SELECT 1 as col0", database=glue_database)
+    assert df.query_metadata["WorkGroup"] == workgroup1