Add catalog_versioning to parameters handling. #342

Author: igorborgest

awswrangler/catalog/_create.py

@@ -163,6 +163,7 @@ def _create_table(  # pylint: disable=too-many-branches,too-many-statements
 def _upsert_table_parameters(
     parameters: Dict[str, str],
     database: str,
+    catalog_versioning: bool,
     catalog_id: Optional[str],
     table_input: Dict[str, Any],
     boto3_session: Optional[boto3.Session],
@@ -180,20 +181,25 @@ def _upsert_table_parameters(
             catalog_id=catalog_id,
             boto3_session=boto3_session,
             table_input=table_input,
+            catalog_versioning=catalog_versioning,
         )
     return pars
 
 
 def _overwrite_table_parameters(
     parameters: Dict[str, str],
     database: str,
+    catalog_versioning: bool,
     catalog_id: Optional[str],
     table_input: Dict[str, Any],
     boto3_session: Optional[boto3.Session],
 ) -> Dict[str, str]:
     table_input["Parameters"] = parameters
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
-    client_glue.update_table(**_catalog_id(catalog_id=catalog_id, DatabaseName=database, TableInput=table_input))
+    skip_archive: bool = not catalog_versioning
+    client_glue.update_table(
+        **_catalog_id(catalog_id=catalog_id, DatabaseName=database, TableInput=table_input, SkipArchive=skip_archive)
+    )
     return parameters
 
 
@@ -346,6 +352,7 @@ def upsert_table_parameters(
     parameters: Dict[str, str],
     database: str,
     table: str,
+    catalog_versioning: bool = False,
     catalog_id: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> Dict[str, str]:
@@ -359,6 +366,8 @@ def upsert_table_parameters(
         Database name.
     table : str
         Table name.
+    catalog_versioning : bool
+        If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it.
     catalog_id : str, optional
         The ID of the Data Catalog from which to retrieve Databases.
         If none is provided, the AWS account ID is used by default.
@@ -386,7 +395,12 @@ def upsert_table_parameters(
     if table_input is None:
         raise exceptions.InvalidArgumentValue(f"Table {database}.{table} does not exist.")
     return _upsert_table_parameters(
-        parameters=parameters, database=database, boto3_session=session, catalog_id=catalog_id, table_input=table_input,
+        parameters=parameters,
+        database=database,
+        boto3_session=session,
+        catalog_id=catalog_id,
+        table_input=table_input,
+        catalog_versioning=catalog_versioning,
     )
 
 
@@ -395,6 +409,7 @@ def overwrite_table_parameters(
     parameters: Dict[str, str],
     database: str,
     table: str,
+    catalog_versioning: bool = False,
     catalog_id: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> Dict[str, str]:
@@ -408,6 +423,8 @@ def overwrite_table_parameters(
         Database name.
     table : str
         Table name.
+    catalog_versioning : bool
+        If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it.
     catalog_id : str, optional
         The ID of the Data Catalog from which to retrieve Databases.
         If none is provided, the AWS account ID is used by default.
@@ -435,7 +452,12 @@ def overwrite_table_parameters(
     if table_input is None:
         raise exceptions.InvalidTable(f"Table {table} does not exist on database {database}.")
     return _overwrite_table_parameters(
-        parameters=parameters, database=database, catalog_id=catalog_id, table_input=table_input, boto3_session=session,
+        parameters=parameters,
+        database=database,
+        catalog_id=catalog_id,
+        table_input=table_input,
+        boto3_session=session,
+        catalog_versioning=catalog_versioning,
     )
 
 

tests/test__routines.py

@@ -27,6 +27,7 @@ def test_routine_0(glue_database, glue_table, path, use_threads, concurrent_part
         use_threads=use_threads,
         concurrent_partitioning=concurrent_partitioning,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     wr.s3.wait_objects_exist(paths=paths, use_threads=use_threads)
     df2 = wr.athena.read_sql_table(glue_table, glue_database, use_threads=use_threads)
     assert df.shape == df2.shape
@@ -55,6 +56,7 @@ def test_routine_0(glue_database, glue_table, path, use_threads, concurrent_part
         use_threads=use_threads,
         concurrent_partitioning=concurrent_partitioning,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     wr.s3.wait_objects_exist(paths=paths, use_threads=use_threads)
     df2 = wr.athena.read_sql_table(glue_table, glue_database, use_threads=use_threads)
     assert df.shape == df2.shape
@@ -83,6 +85,7 @@ def test_routine_0(glue_database, glue_table, path, use_threads, concurrent_part
         use_threads=use_threads,
         concurrent_partitioning=concurrent_partitioning,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     wr.s3.wait_objects_exist(paths=paths, use_threads=use_threads)
     df2 = wr.athena.read_sql_table(glue_table, glue_database, use_threads=use_threads)
     assert len(df.columns) == len(df2.columns)
@@ -112,6 +115,7 @@ def test_routine_0(glue_database, glue_table, path, use_threads, concurrent_part
         use_threads=use_threads,
         concurrent_partitioning=concurrent_partitioning,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     wr.s3.wait_objects_exist(paths=paths, use_threads=use_threads)
     df2 = wr.athena.read_sql_table(glue_table, glue_database, use_threads=use_threads)
     assert len(df2.columns) == 2
@@ -142,6 +146,7 @@ def test_routine_0(glue_database, glue_table, path, use_threads, concurrent_part
         use_threads=use_threads,
         concurrent_partitioning=concurrent_partitioning,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     wr.s3.wait_objects_exist(paths=paths, use_threads=use_threads)
     df2 = wr.athena.read_sql_table(glue_table, glue_database, use_threads=use_threads)
     assert len(df2.columns) == 3
@@ -174,6 +179,7 @@ def test_routine_0(glue_database, glue_table, path, use_threads, concurrent_part
         use_threads=use_threads,
         concurrent_partitioning=concurrent_partitioning,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     wr.s3.wait_objects_exist(paths=paths, use_threads=use_threads)
     df2 = wr.athena.read_sql_table(glue_table, glue_database, use_threads=use_threads)
     assert df.shape == df2.shape
@@ -204,6 +210,7 @@ def test_routine_0(glue_database, glue_table, path, use_threads, concurrent_part
         concurrent_partitioning=concurrent_partitioning,
         use_threads=use_threads,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     wr.s3.wait_objects_exist(paths=paths, use_threads=use_threads)
     df2 = wr.athena.read_sql_table(glue_table, glue_database, use_threads=use_threads)
     assert len(df2.columns) == 2
@@ -235,6 +242,7 @@ def test_routine_0(glue_database, glue_table, path, use_threads, concurrent_part
         use_threads=use_threads,
         concurrent_partitioning=concurrent_partitioning,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     wr.s3.wait_objects_exist(paths=paths, use_threads=use_threads)
     df2 = wr.athena.read_sql_table(glue_table, glue_database, use_threads=use_threads)
     assert len(df2.columns) == 3
@@ -268,6 +276,7 @@ def test_routine_1(glue_database, glue_table, path):
         parameters={"num_cols": str(len(df.columns)), "num_rows": str(len(df.index))},
         columns_comments={"c0": "0"},
     )
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     df2 = wr.athena.read_sql_table(glue_table, glue_database)
     assert df.shape == df2.shape
     assert df.c0.sum() == df2.c0.sum()
@@ -294,6 +303,7 @@ def test_routine_1(glue_database, glue_table, path):
         parameters={"num_cols": str(len(df.columns)), "num_rows": str(len(df.index))},
         columns_comments={"c1": "1"},
     )
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     df2 = wr.athena.read_sql_table(glue_table, glue_database)
     assert df.shape == df2.shape
     assert df.c1.sum() == df2.c1.sum()
@@ -320,6 +330,7 @@ def test_routine_1(glue_database, glue_table, path):
         parameters={"num_cols": str(len(df.columns)), "num_rows": str(len(df.index) * 2)},
         columns_comments={"c1": "1"},
     )
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     df2 = wr.athena.read_sql_table(glue_table, glue_database)
     assert len(df.columns) == len(df2.columns)
     assert len(df.index) * 2 == len(df2.index)
@@ -348,6 +359,7 @@ def test_routine_1(glue_database, glue_table, path):
         parameters={"num_cols": "2", "num_rows": "9"},
         columns_comments={"c1": "1", "c2": "2"},
     )
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     df2 = wr.athena.read_sql_table(glue_table, glue_database)
     assert len(df2.columns) == 2
     assert len(df2.index) == 9
@@ -376,6 +388,7 @@ def test_routine_1(glue_database, glue_table, path):
         parameters={"num_cols": "2", "num_rows": "2"},
         columns_comments={"c0": "zero", "c1": "one"},
     )
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     df2 = wr.athena.read_sql_table(glue_table, glue_database)
     assert df.shape == df2.shape
     assert df.c1.sum() == df2.c1.astype(int).sum()
@@ -405,6 +418,7 @@ def test_routine_1(glue_database, glue_table, path):
         parameters={"num_cols": "2", "num_rows": "3"},
         columns_comments={"c0": "zero", "c1": "one"},
     )
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     df2 = wr.athena.read_sql_table(glue_table, glue_database)
     assert len(df2.columns) == 2
     assert len(df2.index) == 3
@@ -435,6 +449,7 @@ def test_routine_1(glue_database, glue_table, path):
         parameters={"num_cols": "3", "num_rows": "4"},
         columns_comments={"c0": "zero", "c1": "one", "c2": "two"},
     )
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     df2 = wr.athena.read_sql_table(glue_table, glue_database)
     assert len(df2.columns) == 3
     assert len(df2.index) == 4

tests/test_athena.py

@@ -531,6 +531,7 @@ def test_catalog_versioning(path, glue_database, glue_table):
     paths = wr.s3.to_parquet(
         df=df, path=path, dataset=True, database=glue_database, table=glue_table, mode="overwrite"
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 1
     wr.s3.wait_objects_exist(paths=paths, use_threads=False)
     df = wr.athena.read_sql_table(table=glue_table, database=glue_database)
     assert len(df.index) == 2
@@ -548,6 +549,7 @@ def test_catalog_versioning(path, glue_database, glue_table):
         mode="overwrite",
         catalog_versioning=True,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 2
     wr.s3.wait_objects_exist(paths=paths1, use_threads=False)
     df = wr.athena.read_sql_table(table=glue_table, database=glue_database)
     assert len(df.index) == 2
@@ -566,6 +568,7 @@ def test_catalog_versioning(path, glue_database, glue_table):
         catalog_versioning=True,
         index=False,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 3
     wr.s3.wait_objects_exist(paths=paths2, use_threads=False)
     wr.s3.wait_objects_not_exist(paths=paths1, use_threads=False)
     df = wr.athena.read_sql_table(table=glue_table, database=glue_database)
@@ -585,6 +588,7 @@ def test_catalog_versioning(path, glue_database, glue_table):
         catalog_versioning=False,
         index=False,
     )["paths"]
+    assert wr.catalog.get_table_number_of_versions(table=glue_table, database=glue_database) == 3
     wr.s3.wait_objects_exist(paths=paths3, use_threads=False)
     wr.s3.wait_objects_not_exist(paths=paths2, use_threads=False)
     df = wr.athena.read_sql_table(table=glue_table, database=glue_database)