Removing delete_table operations from catalog._create_table() and add catalog_versioning arg. #198

Author: igorborgest

awswrangler/catalog.py

@@ -93,6 +93,7 @@ def create_parquet_table(
     parameters: Optional[Dict[str, str]] = None,
     columns_comments: Optional[Dict[str, str]] = None,
     mode: str = "overwrite",
+    catalog_versioning: bool = False,
     boto3_session: Optional[boto3.Session] = None,
 ) -> None:
     """Create a Parquet Table (Metadata Only) in the AWS Glue Catalog.
@@ -121,6 +122,8 @@ def create_parquet_table(
         Columns names and the related comments (e.g. {'col0': 'Column 0.', 'col1': 'Column 1.', 'col2': 'Partition.'}).
     mode: str
         'overwrite' to recreate any possible existing table or 'append' to keep any possible existing table.
+    catalog_versioning : bool
+        If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
 
@@ -157,6 +160,7 @@ def create_parquet_table(
         parameters=parameters,
         columns_comments=columns_comments,
         mode=mode,
+        catalog_versioning=catalog_versioning,
         boto3_session=boto3_session,
         table_input=table_input,
     )
@@ -865,6 +869,7 @@ def create_csv_table(
     parameters: Optional[Dict[str, str]] = None,
     columns_comments: Optional[Dict[str, str]] = None,
     mode: str = "overwrite",
+    catalog_versioning: bool = False,
     sep: str = ",",
     boto3_session: Optional[boto3.Session] = None,
 ) -> None:
@@ -884,16 +889,18 @@ def create_csv_table(
         Dictionary with keys as column names and vales as data types (e.g. {'col0': 'bigint', 'col1': 'double'}).
     partitions_types: Dict[str, str], optional
         Dictionary with keys as partition names and values as data types (e.g. {'col2': 'date'}).
-    compression: str, optional
+    compression : str, optional
         Compression style (``None``, ``gzip``, etc).
-    description: str, optional
+    description : str, optional
         Table description
-    parameters: Dict[str, str], optional
+    parameters : Dict[str, str], optional
         Key/value pairs to tag the table.
     columns_comments: Dict[str, str], optional
         Columns names and the related comments (e.g. {'col0': 'Column 0.', 'col1': 'Column 1.', 'col2': 'Partition.'}).
-    mode: str
+    mode : str
         'overwrite' to recreate any possible axisting table or 'append' to keep any possible axisting table.
+    catalog_versioning : bool
+        If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it.
     sep : str
         String of length 1. Field delimiter for the output file.
     boto3_session : boto3.Session(), optional
@@ -937,6 +944,7 @@ def create_csv_table(
         parameters=parameters,
         columns_comments=columns_comments,
         mode=mode,
+        catalog_versioning=catalog_versioning,
         boto3_session=boto3_session,
         table_input=table_input,
     )
@@ -949,6 +957,7 @@ def _create_table(
     parameters: Optional[Dict[str, str]],
     columns_comments: Optional[Dict[str, str]],
     mode: str,
+    catalog_versioning: bool,
     boto3_session: Optional[boto3.Session],
     table_input: Dict[str, Any],
 ):
@@ -967,10 +976,14 @@ def _create_table(
             if name in columns_comments:
                 par["Comment"] = columns_comments[name]
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
+    client_glue: boto3.client = _utils.client(service_name="glue", session=session)
     exist: bool = does_table_exist(database=database, table=table, boto3_session=session)
-    if (mode == "overwrite") or (exist is False):
-        delete_table_if_exists(database=database, table=table, boto3_session=session)
-        client_glue: boto3.client = _utils.client(service_name="glue", session=session)
+    if mode not in ("overwrite", "append"):  # pragma: no cover
+        raise exceptions.InvalidArgument(f"{mode} is not a valid mode. It must be 'overwrite' or 'append'.")
+    if (exist is True) and (mode == "overwrite"):
+        skip_archive: bool = not catalog_versioning
+        client_glue.update_table(DatabaseName=database, TableInput=table_input, SkipArchive=skip_archive)
+    elif exist is False:
         client_glue.create_table(DatabaseName=database, TableInput=table_input)
 
 

awswrangler/s3.py

@@ -433,6 +433,7 @@ def to_csv(  # pylint: disable=too-many-arguments
     dataset: bool = False,
     partition_cols: Optional[List[str]] = None,
     mode: Optional[str] = None,
+    catalog_versioning: bool = False,
     database: Optional[str] = None,
     table: Optional[str] = None,
     dtype: Optional[Dict[str, str]] = None,
@@ -483,6 +484,8 @@ def to_csv(  # pylint: disable=too-many-arguments
         List of column names that will be used to create partitions. Only takes effect if dataset=True.
     mode: str, optional
         ``append`` (Default), ``overwrite``, ``overwrite_partitions``. Only takes effect if dataset=True.
+    catalog_versioning : bool
+        If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it.
     database : str, optional
         Glue/Athena catalog: Database name.
     table : str, optional
@@ -677,6 +680,7 @@ def to_csv(  # pylint: disable=too-many-arguments
                     columns_comments=columns_comments,
                     boto3_session=session,
                     mode="overwrite",
+                    catalog_versioning=catalog_versioning,
                     sep=sep,
                 )
             if partitions_values:
@@ -846,6 +850,7 @@ def to_parquet(  # pylint: disable=too-many-arguments
     dataset: bool = False,
     partition_cols: Optional[List[str]] = None,
     mode: Optional[str] = None,
+    catalog_versioning: bool = False,
     database: Optional[str] = None,
     table: Optional[str] = None,
     dtype: Optional[Dict[str, str]] = None,
@@ -893,6 +898,8 @@ def to_parquet(  # pylint: disable=too-many-arguments
         List of column names that will be used to create partitions. Only takes effect if dataset=True.
     mode: str, optional
         ``append`` (Default), ``overwrite``, ``overwrite_partitions``. Only takes effect if dataset=True.
+    catalog_versioning : bool
+        If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it.
     database : str, optional
         Glue/Athena catalog: Database name.
     table : str, optional
@@ -1092,6 +1099,7 @@ def to_parquet(  # pylint: disable=too-many-arguments
                     columns_comments=columns_comments,
                     boto3_session=session,
                     mode="overwrite",
+                    catalog_versioning=catalog_versioning,
                 )
             if partitions_values:
                 _logger.debug("partitions_values:\n%s", partitions_values)
@@ -1838,6 +1846,7 @@ def store_parquet_metadata(
     columns_comments: Optional[Dict[str, str]] = None,
     compression: Optional[str] = None,
     mode: str = "overwrite",
+    catalog_versioning: bool = False,
     boto3_session: Optional[boto3.Session] = None,
 ) -> Tuple[Dict[str, str], Optional[Dict[str, str]], Optional[Dict[str, List[str]]]]:
     """Infer and store parquet metadata on AWS Glue Catalog.
@@ -1879,6 +1888,8 @@ def store_parquet_metadata(
         Compression style (``None``, ``snappy``, ``gzip``, etc).
     mode: str
         'overwrite' to recreate any possible existing table or 'append' to keep any possible existing table.
+    catalog_versioning : bool
+        If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
 
@@ -1924,6 +1935,7 @@ def store_parquet_metadata(
         parameters=parameters,
         columns_comments=columns_comments,
         mode=mode,
+        catalog_versioning=catalog_versioning,
         boto3_session=session,
     )
     partitions_values: Dict[str, List[str]] = _data_types.athena_partitions_from_pyarrow_partitions(

testing/test_awswrangler/test_data_lake.py

@@ -1281,3 +1281,70 @@ def test_athena_nested(bucket, database):
     df2 = wr.athena.read_sql_query(sql=f"SELECT c0, c1, c2, c4 FROM {table}", database=database)
     assert len(df2.index) == 2
     assert len(df2.columns) == 4
+
+
+def test_catalog_versioning(bucket, database):
+    table = "test_catalog_versioning"
+    wr.catalog.delete_table_if_exists(database=database, table=table)
+    path = f"s3://{bucket}/{table}/"
+    wr.s3.delete_objects(path=path)
+
+    # Version 0
+    df = pd.DataFrame({"c0": [1, 2]})
+    paths = wr.s3.to_parquet(df=df, path=path, dataset=True, database=database, table=table, mode="overwrite")["paths"]
+    wr.s3.wait_objects_exist(paths=paths, use_threads=False)
+    df = wr.athena.read_sql_table(table=table, database=database)
+    assert len(df.index) == 2
+    assert len(df.columns) == 1
+    assert str(df.c0.dtype).startswith("Int")
+
+    # Version 1
+    df = pd.DataFrame({"c1": ["foo", "boo"]})
+    paths = wr.s3.to_parquet(
+        df=df, path=path, dataset=True, database=database, table=table, mode="overwrite", catalog_versioning=True
+    )["paths"]
+    wr.s3.wait_objects_exist(paths=paths, use_threads=False)
+    df = wr.athena.read_sql_table(table=table, database=database)
+    assert len(df.index) == 2
+    assert len(df.columns) == 1
+    assert str(df.c1.dtype) == "string"
+
+    # Version 2
+    df = pd.DataFrame({"c1": [1.0, 2.0]})
+    paths = wr.s3.to_csv(
+        df=df,
+        path=path,
+        dataset=True,
+        database=database,
+        table=table,
+        mode="overwrite",
+        catalog_versioning=True,
+        index=False,
+    )["paths"]
+    wr.s3.wait_objects_exist(paths=paths, use_threads=False)
+    df = wr.athena.read_sql_table(table=table, database=database)
+    assert len(df.index) == 2
+    assert len(df.columns) == 1
+    assert str(df.c1.dtype).startswith("float")
+
+    # Version 3 (removing version 2)
+    df = pd.DataFrame({"c1": [True, False]})
+    paths = wr.s3.to_csv(
+        df=df,
+        path=path,
+        dataset=True,
+        database=database,
+        table=table,
+        mode="overwrite",
+        catalog_versioning=False,
+        index=False,
+    )["paths"]
+    wr.s3.wait_objects_exist(paths=paths, use_threads=False)
+    df = wr.athena.read_sql_table(table=table, database=database)
+    assert len(df.index) == 2
+    assert len(df.columns) == 1
+    assert str(df.c1.dtype).startswith("boolean")
+
+    # Cleaning Up
+    wr.catalog.delete_table_if_exists(database=database, table=table)
+    wr.s3.delete_objects(path=path)