diff --git a/awswrangler/athena/_executions.py b/awswrangler/athena/_executions.py index b612d6889..e10080824 100644 --- a/awswrangler/athena/_executions.py +++ b/awswrangler/athena/_executions.py @@ -35,7 +35,7 @@ def start_query_execution( sql: str, database: Optional[str] = None, s3_output: Optional[str] = None, - workgroup: Optional[str] = None, + workgroup: str = "primary", encryption: Optional[str] = None, kms_key: Optional[str] = None, params: Union[Dict[str, Any], List[str], None] = None, @@ -62,8 +62,8 @@ def start_query_execution( AWS Glue/Athena database name. s3_output : str, optional AWS S3 path. - workgroup : str, optional - Athena workgroup. + workgroup : str + Athena workgroup. Primary by default. encryption : str, optional None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'. kms_key : str, optional diff --git a/awswrangler/athena/_executions.pyi b/awswrangler/athena/_executions.pyi index 142cedb74..fbd68284c 100644 --- a/awswrangler/athena/_executions.pyi +++ b/awswrangler/athena/_executions.pyi @@ -17,7 +17,7 @@ def start_query_execution( sql: str, database: Optional[str] = ..., s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., params: Union[Dict[str, Any], List[str], None] = ..., @@ -34,7 +34,7 @@ def start_query_execution( *, database: Optional[str] = ..., s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., params: Union[Dict[str, Any], List[str], None] = ..., @@ -51,7 +51,7 @@ def start_query_execution( *, database: Optional[str] = ..., s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., params: Union[Dict[str, Any], List[str], None] = ..., diff --git a/awswrangler/athena/_read.py b/awswrangler/athena/_read.py index b0f258c26..eecfdc5a1 100644 --- a/awswrangler/athena/_read.py +++ b/awswrangler/athena/_read.py @@ -774,7 +774,7 @@ def read_sql_query( # pylint: disable=too-many-arguments,too-many-locals categories: Optional[List[str]] = None, chunksize: Optional[Union[int, bool]] = None, s3_output: Optional[str] = None, - workgroup: Optional[str] = None, + workgroup: str = "primary", encryption: Optional[str] = None, kms_key: Optional[str] = None, keep_files: bool = True, @@ -929,8 +929,8 @@ def read_sql_query( # pylint: disable=too-many-arguments,too-many-locals If an `INTEGER` is passed awswrangler will iterate on the data by number of rows equal the received INTEGER. s3_output : str, optional Amazon S3 path. - workgroup : str, optional - Athena workgroup. + workgroup : str + Athena workgroup. Primary by default. encryption : str, optional Valid values: [None, 'SSE_S3', 'SSE_KMS']. Notice: 'CSE_KMS' is not supported. kms_key : str, optional @@ -1122,7 +1122,7 @@ def read_sql_table( categories: Optional[List[str]] = None, chunksize: Optional[Union[int, bool]] = None, s3_output: Optional[str] = None, - workgroup: Optional[str] = None, + workgroup: str = "primary", encryption: Optional[str] = None, kms_key: Optional[str] = None, keep_files: bool = True, @@ -1272,8 +1272,8 @@ def read_sql_table( If an `INTEGER` is passed awswrangler will iterate on the data by number of rows equal the received INTEGER. s3_output : str, optional AWS S3 path. - workgroup : str, optional - Athena workgroup. + workgroup : str + Athena workgroup. Primary by default. encryption : str, optional Valid values: [None, 'SSE_S3', 'SSE_KMS']. Notice: 'CSE_KMS' is not supported. kms_key : str, optional @@ -1364,7 +1364,7 @@ def unload( compression: Optional[str] = None, field_delimiter: Optional[str] = None, partitioned_by: Optional[List[str]] = None, - workgroup: Optional[str] = None, + workgroup: str = "primary", encryption: Optional[str] = None, kms_key: Optional[str] = None, boto3_session: Optional[boto3.Session] = None, @@ -1397,8 +1397,8 @@ def unload( A single-character field delimiter for files in CSV, TSV, and other text formats. partitioned_by : Optional[List[str]] An array list of columns by which the output is partitioned. - workgroup : str, optional - Athena workgroup. + workgroup : str + Athena workgroup. Primary by default. encryption : str, optional Valid values: [None, 'SSE_S3', 'SSE_KMS']. Notice: 'CSE_KMS' is not supported. kms_key : str, optional diff --git a/awswrangler/athena/_read.pyi b/awswrangler/athena/_read.pyi index 459d35fb7..93acb7045 100644 --- a/awswrangler/athena/_read.pyi +++ b/awswrangler/athena/_read.pyi @@ -64,7 +64,7 @@ def read_sql_query( # pylint: disable=too-many-arguments categories: Optional[List[str]] = ..., chunksize: Union[None, Literal[False]] = ..., s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., keep_files: bool = ..., @@ -91,7 +91,7 @@ def read_sql_query( categories: Optional[List[str]] = ..., chunksize: Literal[True], s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., keep_files: bool = ..., @@ -118,7 +118,7 @@ def read_sql_query( categories: Optional[List[str]] = ..., chunksize: bool, s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., keep_files: bool = ..., @@ -145,7 +145,7 @@ def read_sql_query( categories: Optional[List[str]] = ..., chunksize: int, s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., keep_files: bool = ..., @@ -172,7 +172,7 @@ def read_sql_query( categories: Optional[List[str]] = ..., chunksize: Optional[Union[int, bool]], s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., keep_files: bool = ..., @@ -199,7 +199,7 @@ def read_sql_table( categories: Optional[List[str]] = ..., chunksize: Union[None, Literal[False]] = ..., s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., keep_files: bool = ..., @@ -223,7 +223,7 @@ def read_sql_table( categories: Optional[List[str]] = ..., chunksize: Literal[True], s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., keep_files: bool = ..., @@ -247,7 +247,7 @@ def read_sql_table( categories: Optional[List[str]] = ..., chunksize: bool, s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., keep_files: bool = ..., @@ -271,7 +271,7 @@ def read_sql_table( categories: Optional[List[str]] = ..., chunksize: int, s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., keep_files: bool = ..., @@ -295,7 +295,7 @@ def read_sql_table( categories: Optional[List[str]] = ..., chunksize: Optional[Union[int, bool]], s3_output: Optional[str] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., keep_files: bool = ..., @@ -315,7 +315,7 @@ def unload( compression: Optional[str] = ..., field_delimiter: Optional[str] = ..., partitioned_by: Optional[List[str]] = ..., - workgroup: Optional[str] = ..., + workgroup: str = ..., encryption: Optional[str] = ..., kms_key: Optional[str] = ..., boto3_session: Optional[boto3.Session] = ..., diff --git a/awswrangler/athena/_statements.py b/awswrangler/athena/_statements.py index 9c5c50212..c6a043894 100644 --- a/awswrangler/athena/_statements.py +++ b/awswrangler/athena/_statements.py @@ -35,7 +35,7 @@ def _does_statement_exist( def create_prepared_statement( sql: str, statement_name: str, - workgroup: Optional[str] = None, + workgroup: str = "primary", mode: Literal["update", "error"] = "update", boto3_session: Optional[boto3.Session] = None, ) -> None: @@ -50,8 +50,8 @@ def create_prepared_statement( The query string for the prepared statement. statement_name : str The name of the prepared statement. - workgroup : str, optional - The name of the workgroup to which the prepared statement belongs. + workgroup : str + The name of the workgroup to which the prepared statement belongs. Primary by default. mode: str Determines the behaviour if the prepared statement already exists: @@ -72,7 +72,6 @@ def create_prepared_statement( raise exceptions.InvalidArgumentValue("`mode` must be one of 'update' or 'error'.") athena_client = _utils.client("athena", session=boto3_session) - workgroup = workgroup if workgroup else "primary" already_exists = _does_statement_exist(statement_name, workgroup, athena_client) if already_exists and mode == "error": @@ -95,16 +94,14 @@ def create_prepared_statement( @apply_configs -def list_prepared_statements( - workgroup: Optional[str] = None, boto3_session: Optional[boto3.Session] = None -) -> List[str]: +def list_prepared_statements(workgroup: str = "primary", boto3_session: Optional[boto3.Session] = None) -> List[str]: """ List the prepared statements in the specified workgroup. Parameters ---------- - workgroup: str, optional - The name of the workgroup to which the prepared statement belongs. + workgroup: str + The name of the workgroup to which the prepared statement belongs. Primary by default. boto3_session : boto3.Session(), optional Boto3 Session. The default boto3 session will be used if boto3_session receive None. @@ -115,7 +112,6 @@ def list_prepared_statements( Each item is a dictionary with the keys ``StatementName`` and ``LastModifiedTime``. """ athena_client = _utils.client("athena", session=boto3_session) - workgroup = workgroup if workgroup else "primary" response = athena_client.list_prepared_statements(WorkGroup=workgroup) statements = response["PreparedStatements"] @@ -130,7 +126,7 @@ def list_prepared_statements( @apply_configs def delete_prepared_statement( statement_name: str, - workgroup: Optional[str] = None, + workgroup: str = "primary", boto3_session: Optional[boto3.Session] = None, ) -> None: """ @@ -143,7 +139,7 @@ def delete_prepared_statement( statement_name : str The name of the prepared statement. workgroup : str, optional - The name of the workgroup to which the prepared statement belongs. + The name of the workgroup to which the prepared statement belongs. Primary by default. boto3_session : boto3.Session(), optional Boto3 Session. The default boto3 session will be used if boto3_session receive None. diff --git a/awswrangler/athena/_utils.py b/awswrangler/athena/_utils.py index 906bdf2ac..48e57bf9f 100644 --- a/awswrangler/athena/_utils.py +++ b/awswrangler/athena/_utils.py @@ -136,7 +136,7 @@ def _start_query_execution( return response["QueryExecutionId"] -def _get_workgroup_config(session: Optional[boto3.Session] = None, workgroup: Optional[str] = None) -> _WorkGroupConfig: +def _get_workgroup_config(session: Optional[boto3.Session] = None, workgroup: str = "primary") -> _WorkGroupConfig: enforced: bool wg_s3_output: Optional[str] wg_encryption: Optional[str] @@ -472,7 +472,7 @@ def repair_table( database: Optional[str] = None, data_source: Optional[str] = None, s3_output: Optional[str] = None, - workgroup: Optional[str] = None, + workgroup: str = "primary", encryption: Optional[str] = None, kms_key: Optional[str] = None, athena_query_wait_polling_delay: float = _QUERY_WAIT_POLLING_DELAY, @@ -501,8 +501,8 @@ def repair_table( Data Source / Catalog name. If None, 'AwsDataCatalog' is used. s3_output : str, optional AWS S3 path. - workgroup : str, optional - Athena workgroup. + workgroup : str + Athena workgroup. Primary by default. encryption : str, optional None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'. kms_key : str, optional @@ -552,7 +552,7 @@ def describe_table( table: str, database: Optional[str] = None, s3_output: Optional[str] = None, - workgroup: Optional[str] = None, + workgroup: str = "primary", encryption: Optional[str] = None, kms_key: Optional[str] = None, athena_query_wait_polling_delay: float = _QUERY_WAIT_POLLING_DELAY, @@ -577,8 +577,8 @@ def describe_table( AWS Glue/Athena database name. s3_output : str, optional AWS S3 path. - workgroup : str, optional - Athena workgroup. + workgroup : str + Athena workgroup. Primary by default. encryption : str, optional None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'. kms_key : str, optional @@ -641,7 +641,7 @@ def create_ctas_table( # pylint: disable=too-many-locals bucketing_info: Optional[typing.BucketingInfoTuple] = None, field_delimiter: Optional[str] = None, schema_only: bool = False, - workgroup: Optional[str] = None, + workgroup: str = "primary", data_source: Optional[str] = None, encryption: Optional[str] = None, kms_key: Optional[str] = None, @@ -686,8 +686,8 @@ def create_ctas_table( # pylint: disable=too-many-locals The single-character field delimiter for files in CSV, TSV, and text files. schema_only : bool, optional _description_, by default False - workgroup : Optional[str], optional - Athena workgroup. + workgroup : str + Athena workgroup. Primary by default. data_source : Optional[str], optional Data Source / Catalog name. If None, 'AwsDataCatalog' is used. encryption : str, optional @@ -856,7 +856,7 @@ def show_create_table( table: str, database: Optional[str] = None, s3_output: Optional[str] = None, - workgroup: Optional[str] = None, + workgroup: str = "primary", encryption: Optional[str] = None, kms_key: Optional[str] = None, athena_query_wait_polling_delay: float = _QUERY_WAIT_POLLING_DELAY, @@ -880,8 +880,8 @@ def show_create_table( AWS Glue/Athena database name. s3_output : str, optional AWS S3 path. - workgroup : str, optional - Athena workgroup. + workgroup : str + Athena workgroup. Primary by default. encryption : str, optional None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'. kms_key : str, optional diff --git a/awswrangler/athena/_write_iceberg.py b/awswrangler/athena/_write_iceberg.py index a1ac1c5f5..7e417dc20 100644 --- a/awswrangler/athena/_write_iceberg.py +++ b/awswrangler/athena/_write_iceberg.py @@ -196,7 +196,7 @@ def to_iceberg( partition_cols: Optional[List[str]] = None, keep_files: bool = True, data_source: Optional[str] = None, - workgroup: Optional[str] = None, + workgroup: str = "primary", encryption: Optional[str] = None, kms_key: Optional[str] = None, boto3_session: Optional[boto3.Session] = None, @@ -237,8 +237,8 @@ def to_iceberg( Whether staging files produced by Athena are retained. 'True' by default. data_source : str, optional Data Source / Catalog name. If None, 'AwsDataCatalog' will be used by default. - workgroup : str, optional - Athena workgroup. + workgroup : str + Athena workgroup. Primary by default. encryption : str, optional Valid values: [None, 'SSE_S3', 'SSE_KMS']. Notice: 'CSE_KMS' is not supported. kms_key : str, optional diff --git a/awswrangler/neptune/_neptune.py b/awswrangler/neptune/_neptune.py index d13b17760..d8fbd8b03 100644 --- a/awswrangler/neptune/_neptune.py +++ b/awswrangler/neptune/_neptune.py @@ -176,6 +176,8 @@ def to_property_graph( raise exceptions.InvalidArgumentValue( "DataFrame must contain at least a ~id and a ~label column to be saved to Amazon Neptune" ) + if df.empty: + raise exceptions.EmptyDataFrame("DataFrame cannot be empty.") # Loop through items in the DF for index, row in df.iterrows(): diff --git a/tests/unit/test_neptune.py b/tests/unit/test_neptune.py index 012bbbda0..28f2581d6 100644 --- a/tests/unit/test_neptune.py +++ b/tests/unit/test_neptune.py @@ -492,6 +492,13 @@ def test_sparql_write_quads(neptune_endpoint, neptune_port) -> Dict[str, Any]: assert len(batch_df.index) == len(final_df.index) + 50 +def test_gremlin_write_empty(neptune_endpoint, neptune_port) -> Dict[str, Any]: + client = wr.neptune.connect(neptune_endpoint, neptune_port) + + with pytest.raises(wr.exceptions.EmptyDataFrame): + wr.neptune.to_property_graph(client, pd.DataFrame(columns=["~id", "~label"])) + + def _create_dummy_vertex(label: str = "foo") -> Dict[str, Any]: return { "~id": str(uuid.uuid4()),