Fix missing catalog_id arguments. #399

Author: igorborgest

awswrangler/catalog/_add.py

@@ -39,6 +39,7 @@ def add_csv_partitions(
     database: str,
     table: str,
     partitions_values: Dict[str, List[str]],
+    catalog_id: Optional[str] = None,
     compression: Optional[str] = None,
     sep: str = ",",
     boto3_session: Optional[boto3.Session] = None,
@@ -54,6 +55,9 @@ def add_csv_partitions(
     partitions_values: Dict[str, List[str]]
         Dictionary with keys as S3 path locations and values as a list of partitions values as str
         (e.g. {'s3://bucket/prefix/y=2020/m=10/': ['2020', '10']}).
+    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.
     compression: str, optional
         Compression style (``None``, ``gzip``, etc).
     sep : str
@@ -84,7 +88,7 @@ def add_csv_partitions(
         _csv_partition_definition(location=k, values=v, compression=compression, sep=sep)
         for k, v in partitions_values.items()
     ]
-    _add_partitions(database=database, table=table, boto3_session=boto3_session, inputs=inputs)
+    _add_partitions(database=database, table=table, boto3_session=boto3_session, inputs=inputs, catalog_id=catalog_id)
 
 
 @apply_configs

awswrangler/s3/_write_parquet.py

@@ -561,6 +561,7 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals
                     partitions_values=partitions_values,
                     compression=compression,
                     boto3_session=session,
+                    catalog_id=catalog_id,
                 )
     return {"paths": paths, "partitions_values": partitions_values}
 

awswrangler/s3/_write_text.py

@@ -436,7 +436,12 @@ def to_csv(  # pylint: disable=too-many-arguments,too-many-locals
             if partitions_values and (regular_partitions is True):
                 _logger.debug("partitions_values:\n%s", partitions_values)
                 catalog.add_csv_partitions(
-                    database=database, table=table, partitions_values=partitions_values, boto3_session=session, sep=sep
+                    database=database,
+                    table=table,
+                    partitions_values=partitions_values,
+                    boto3_session=session,
+                    sep=sep,
+                    catalog_id=catalog_id,
                 )
     return {"paths": paths, "partitions_values": partitions_values}