Feature: Glue Governed Tables (#560) (#571)

* Feature: Write to Glue Governed Tables (#560)

* Initial Commit

* Minor - Refactoring Work Units Logic

* Major - Checkpoint w/ functional read code/example

* Initial Commit

* Minor - Refactoring Work Units Logic

* Major - Checkpoint w/ functional read code/example

* Minor - Removing unnecessary ensure_session

* Minor - Adding changes from comments and review

* Minor - Adding Abort, Begin, Commit and Extend transactions

* Minor - Adding missing functions

* Minor - Adding missing @Property

* Minor - Disable too many public methods

* Minor - Checkpoint

* Major - Governed tables write operations tested

* Minor - Adding validate flow on branches

* Minor - reducing static checks

* Minor - Adding to_csv code

* Minor - Disabling too-many-branches

* Major - Ready for release

* Minor - Proofreading

* Minor - Removing needless use_threads argument

* Minor - Removing the need to specify table_type when table is already created

* Minor - Fixing _catalog_id call

* Minor - Clarifying SQL filter operation

* Minor - Removing type ignore

* Minor - Reducing scope gitworkflow

* Minor - Fixing _sanitize_name

* Minor - Adding map_types flag

* Minor - Aligning optional path argument with main branch

* Minor tests adjustments

* Minor - Removing Chunked parameter

* Fixing issues from diverged branch

* Major - M1 Launch (Stable)

* Mmproving read by concatenating zero copies of arrow ttables

* [skip ci] - Minor - Killing thread

* [skip ci] - Minor - Passing client instead of session

* Major - Adding Metadata Transaction API changes

* Minor - Adding query as of time to test

* [skip ci] - syncing with main branch

* Merge main and adapt tests to API changes from Erie team

* Merge main

* Merging main - move to poetry

* Merge main and resolve conflicts

* Minor - Sync with main

* Merging with main

* Lint

* Fixing get_table_obj retries

* Green tests

* Minor -  Fixing automated merge

* LakeFormation test infra

* Commit protocol change - Erie

* [skip ci] - Minor - Fixing catalog unit test

* [skip ci] - Minor - Adding transaction_id to does_table_exist

* Minor - Missing projection_storage_location_template

* Upgrading botocore

* xfail moto

* Adding s3fs to tox

* LF concurrent modification exception

* catalog.py test

* lint

Author: jaidisido

.github/workflows/minimal-tests.yml

@@ -4,7 +4,6 @@ on:
   push:
     branches:
       - main
-      - main-governed-tables
   pull_request:
     branches:
       - main

.github/workflows/static-checking.yml

@@ -4,7 +4,6 @@ on:
   push:
     branches:
       - main
-      - main-governed-tables
   pull_request:
     branches:
       - main

CONTRIBUTING.md

@@ -153,9 +153,9 @@ or
 
 ``cd scripts``
 
-* Deploy the Cloudformation template `base.yaml`
+* Deploy the `base` CDK stack
 
-``./deploy-base.sh``
+``./deploy-stack.sh base``
 
 * Return to the project root directory
 
@@ -175,7 +175,7 @@ or
 
 * [OPTIONAL] To remove the base test environment cloud formation stack post testing:
 
-``./test_infra/scripts/delete-base.sh``
+``./test_infra/scripts/delete-stack.sh base``
 
 ### Full test environment
 
@@ -210,14 +210,18 @@ or
 
 ``cd scripts``
 
-* Deploy the Cloudformation templates `base.yaml` and `databases.yaml`. This step could take about 15 minutes to deploy.
+* Deploy the `base` and `databases` CDK stacks. This step could take about 15 minutes to deploy.
 
-``./deploy-base.sh``
-``./deploy-databases.sh``
+``./deploy-stack.sh base``
+``./deploy-stack.sh databases``
 
-* [OPTIONAL] Deploy the Cloudformation template `opensearch.yaml` (if you need to test Amazon OpenSearch Service). This step could take about 15 minutes to deploy.
+* [OPTIONAL] Deploy the `lakeformation` CDK stack (if you need to test against the AWS Lake Formation Service). You must ensure Lake Formation is enabled in the account.
 
-``./deploy-opensearch.sh``
+``./deploy-stack.sh lakeformation``
+
+* [OPTIONAL] Deploy the `opensearch` CDK stack (if you need to test against the Amazon OpenSearch Service). This step could take about 15 minutes to deploy.
+
+``./deploy-stack.sh opensearch``
 
 * Go to the `EC2 -> SecurityGroups` console, open the `aws-data-wrangler-*` security group and configure to accept your IP from any TCP port.
   - Alternatively run:
@@ -254,9 +258,9 @@ or
 
 * [OPTIONAL] To remove the base test environment cloud formation stack post testing:
 
-``./test_infra/scripts/delete-base.sh``
+``./test_infra/scripts/delete-stack.sh base``
 
-``./test_infra/scripts/delete-databases.sh``
+``./test_infra/scripts/delete-stack.sh databases``
 
 ## Recommended Visual Studio Code Recommended setting
 

README.md

@@ -139,6 +139,7 @@ FROM "sampleDB"."sampleTable" ORDER BY time DESC LIMIT 3
   - [029 - S3 Select](https://github.com/awslabs/aws-data-wrangler/blob/main/tutorials/029%20-%20S3%20Select.ipynb)
   - [030 - Data Api](https://github.com/awslabs/aws-data-wrangler/blob/main/tutorials/030%20-%20Data%20Api.ipynb)
   - [031 - OpenSearch](https://github.com/awslabs/aws-data-wrangler/blob/main/tutorials/031%20-%20OpenSearch.ipynb)
+  - [032 - Lake Formation Governed Tables](https://github.com/awslabs/aws-data-wrangler/blob/main/tutorials/032%20-%Lake%20Formation%20Governed%20Tables.ipynb)
 - [**API Reference**](https://aws-data-wrangler.readthedocs.io/en/2.12.1/api.html)
   - [Amazon S3](https://aws-data-wrangler.readthedocs.io/en/2.12.1/api.html#amazon-s3)
   - [AWS Glue Catalog](https://aws-data-wrangler.readthedocs.io/en/2.12.1/api.html#aws-glue-catalog)

awswrangler/__init__.py

@@ -16,6 +16,7 @@
     dynamodb,
     emr,
     exceptions,
+    lakeformation,
     mysql,
     opensearch,
     postgresql,
@@ -44,6 +45,7 @@
     "s3",
     "sts",
     "redshift",
+    "lakeformation",
     "mysql",
     "postgresql",
     "secretsmanager",

awswrangler/_config.py

@@ -43,14 +43,15 @@ class _ConfigArg(NamedTuple):
     "redshift_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
     "kms_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
     "emr_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
+    "lakeformation_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
     "dynamodb_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
     "secretsmanager_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
     # Botocore config
     "botocore_config": _ConfigArg(dtype=botocore.config.Config, nullable=True),
 }
 
 
-class _Config:  # pylint: disable=too-many-instance-attributes, too-many-public-methods
+class _Config:  # pylint: disable=too-many-instance-attributes,too-many-public-methods
     """Wrangler's Configuration class."""
 
     def __init__(self) -> None:
@@ -63,6 +64,7 @@ def __init__(self) -> None:
         self.redshift_endpoint_url = None
         self.kms_endpoint_url = None
         self.emr_endpoint_url = None
+        self.lakeformation_endpoint_url = None
         self.dynamodb_endpoint_url = None
         self.secretsmanager_endpoint_url = None
         self.botocore_config = None
@@ -356,6 +358,15 @@ def emr_endpoint_url(self) -> Optional[str]:
     def emr_endpoint_url(self, value: Optional[str]) -> None:
         self._set_config_value(key="emr_endpoint_url", value=value)
 
+    @property
+    def lakeformation_endpoint_url(self) -> Optional[str]:
+        """Property lakeformation_endpoint_url."""
+        return cast(Optional[str], self["lakeformation_endpoint_url"])
+
+    @lakeformation_endpoint_url.setter
+    def lakeformation_endpoint_url(self, value: Optional[str]) -> None:
+        self._set_config_value(key="lakeformation_endpoint_url", value=value)
+
     @property
     def dynamodb_endpoint_url(self) -> Optional[str]:
         """Property dynamodb_endpoint_url."""

awswrangler/_utils.py

@@ -93,6 +93,8 @@ def _get_endpoint_url(service_name: str) -> Optional[str]:
         endpoint_url = _config.config.kms_endpoint_url
     elif service_name == "emr" and _config.config.emr_endpoint_url is not None:
         endpoint_url = _config.config.emr_endpoint_url
+    elif service_name == "lakeformation" and _config.config.lakeformation_endpoint_url is not None:
+        endpoint_url = _config.config.lakeformation_endpoint_url
     elif service_name == "dynamodb" and _config.config.dynamodb_endpoint_url is not None:
         endpoint_url = _config.config.dynamodb_endpoint_url
     elif service_name == "secretsmanager" and _config.config.secretsmanager_endpoint_url is not None:

awswrangler/catalog/_add.py

@@ -14,7 +14,7 @@
     _parquet_partition_definition,
     _update_table_definition,
 )
-from awswrangler.catalog._utils import _catalog_id, sanitize_table_name
+from awswrangler.catalog._utils import _catalog_id, _transaction_id, sanitize_table_name
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -300,7 +300,8 @@ def add_column(
     table: str,
     column_name: str,
     column_type: str = "string",
-    column_comment: Optional[str] = "",
+    column_comment: Optional[str] = None,
+    transaction_id: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
     catalog_id: Optional[str] = None,
 ) -> None:
@@ -318,6 +319,8 @@ def add_column(
         Column type.
     column_comment : str
         Column Comment
+    transaction_id: str, optional
+        The ID of the transaction (i.e. used with GOVERNED tables).
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
     catalog_id : str, optional
@@ -341,13 +344,21 @@ def add_column(
     """
     if _check_column_type(column_type):
         client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
-        table_res: Dict[str, Any] = client_glue.get_table(DatabaseName=database, Name=table)
+        table_res: Dict[str, Any] = client_glue.get_table(
+            **_catalog_id(
+                catalog_id=catalog_id,
+                **_transaction_id(transaction_id=transaction_id, DatabaseName=database, Name=table),
+            )
+        )
         table_input: Dict[str, Any] = _update_table_definition(table_res)
         table_input["StorageDescriptor"]["Columns"].append(
             {"Name": column_name, "Type": column_type, "Comment": column_comment}
         )
         res: Dict[str, Any] = client_glue.update_table(
-            **_catalog_id(catalog_id=catalog_id, DatabaseName=database, TableInput=table_input)
+            **_catalog_id(
+                catalog_id=catalog_id,
+                **_transaction_id(transaction_id=transaction_id, DatabaseName=database, TableInput=table_input),
+            )
         )
         if ("Errors" in res) and res["Errors"]:
             for error in res["Errors"]: