Skip to content

Commit 28880d5

Browse files
kukushkingkukushking
authored
fix: Athena/Neptune minor fixes (#2526)
* fix: Raise EmptyDataFrame exception in wr.neptune.to_property_graph * fix: Explicitly use primary Athena workgroup
1 parent d1f525c commit 28880d5

File tree

9 files changed

+59
-54
lines changed

9 files changed

+59
-54
lines changed

‎awswrangler/athena/_executions.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def start_query_execution(
3535
sql: str,
3636
database: Optional[str] = None,
3737
s3_output: Optional[str] = None,
38-
workgroup: Optional[str] = None,
38+
workgroup: str = "primary",
3939
encryption: Optional[str] = None,
4040
kms_key: Optional[str] = None,
4141
params: Union[Dict[str, Any], List[str], None] = None,
@@ -62,8 +62,8 @@ def start_query_execution(
6262
AWS Glue/Athena database name.
6363
s3_output : str, optional
6464
AWS S3 path.
65-
workgroup : str, optional
66-
Athena workgroup.
65+
workgroup : str
66+
Athena workgroup. Primary by default.
6767
encryption : str, optional
6868
None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'.
6969
kms_key : str, optional

‎awswrangler/athena/_executions.pyi‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ def start_query_execution(
1717
sql: str,
1818
database: Optional[str] = ...,
1919
s3_output: Optional[str] = ...,
20-
workgroup: Optional[str] = ...,
20+
workgroup: str = ...,
2121
encryption: Optional[str] = ...,
2222
kms_key: Optional[str] = ...,
2323
params: Union[Dict[str, Any], List[str], None] = ...,
@@ -34,7 +34,7 @@ def start_query_execution(
3434
*,
3535
database: Optional[str] = ...,
3636
s3_output: Optional[str] = ...,
37-
workgroup: Optional[str] = ...,
37+
workgroup: str = ...,
3838
encryption: Optional[str] = ...,
3939
kms_key: Optional[str] = ...,
4040
params: Union[Dict[str, Any], List[str], None] = ...,
@@ -51,7 +51,7 @@ def start_query_execution(
5151
*,
5252
database: Optional[str] = ...,
5353
s3_output: Optional[str] = ...,
54-
workgroup: Optional[str] = ...,
54+
workgroup: str = ...,
5555
encryption: Optional[str] = ...,
5656
kms_key: Optional[str] = ...,
5757
params: Union[Dict[str, Any], List[str], None] = ...,

‎awswrangler/athena/_read.py‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -774,7 +774,7 @@ def read_sql_query( # pylint: disable=too-many-arguments,too-many-locals
774774
categories: Optional[List[str]] = None,
775775
chunksize: Optional[Union[int, bool]] = None,
776776
s3_output: Optional[str] = None,
777-
workgroup: Optional[str] = None,
777+
workgroup: str = "primary",
778778
encryption: Optional[str] = None,
779779
kms_key: Optional[str] = None,
780780
keep_files: bool = True,
@@ -929,8 +929,8 @@ def read_sql_query( # pylint: disable=too-many-arguments,too-many-locals
929929
If an `INTEGER` is passed awswrangler will iterate on the data by number of rows equal the received INTEGER.
930930
s3_output : str, optional
931931
Amazon S3 path.
932-
workgroup : str, optional
933-
Athena workgroup.
932+
workgroup : str
933+
Athena workgroup. Primary by default.
934934
encryption : str, optional
935935
Valid values: [None, 'SSE_S3', 'SSE_KMS']. Notice: 'CSE_KMS' is not supported.
936936
kms_key : str, optional
@@ -1122,7 +1122,7 @@ def read_sql_table(
11221122
categories: Optional[List[str]] = None,
11231123
chunksize: Optional[Union[int, bool]] = None,
11241124
s3_output: Optional[str] = None,
1125-
workgroup: Optional[str] = None,
1125+
workgroup: str = "primary",
11261126
encryption: Optional[str] = None,
11271127
kms_key: Optional[str] = None,
11281128
keep_files: bool = True,
@@ -1272,8 +1272,8 @@ def read_sql_table(
12721272
If an `INTEGER` is passed awswrangler will iterate on the data by number of rows equal the received INTEGER.
12731273
s3_output : str, optional
12741274
AWS S3 path.
1275-
workgroup : str, optional
1276-
Athena workgroup.
1275+
workgroup : str
1276+
Athena workgroup. Primary by default.
12771277
encryption : str, optional
12781278
Valid values: [None, 'SSE_S3', 'SSE_KMS']. Notice: 'CSE_KMS' is not supported.
12791279
kms_key : str, optional
@@ -1364,7 +1364,7 @@ def unload(
13641364
compression: Optional[str] = None,
13651365
field_delimiter: Optional[str] = None,
13661366
partitioned_by: Optional[List[str]] = None,
1367-
workgroup: Optional[str] = None,
1367+
workgroup: str = "primary",
13681368
encryption: Optional[str] = None,
13691369
kms_key: Optional[str] = None,
13701370
boto3_session: Optional[boto3.Session] = None,
@@ -1397,8 +1397,8 @@ def unload(
13971397
A single-character field delimiter for files in CSV, TSV, and other text formats.
13981398
partitioned_by : Optional[List[str]]
13991399
An array list of columns by which the output is partitioned.
1400-
workgroup : str, optional
1401-
Athena workgroup.
1400+
workgroup : str
1401+
Athena workgroup. Primary by default.
14021402
encryption : str, optional
14031403
Valid values: [None, 'SSE_S3', 'SSE_KMS']. Notice: 'CSE_KMS' is not supported.
14041404
kms_key : str, optional

‎awswrangler/athena/_read.pyi‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ def read_sql_query( # pylint: disable=too-many-arguments
6464
categories: Optional[List[str]] = ...,
6565
chunksize: Union[None, Literal[False]] = ...,
6666
s3_output: Optional[str] = ...,
67-
workgroup: Optional[str] = ...,
67+
workgroup: str = ...,
6868
encryption: Optional[str] = ...,
6969
kms_key: Optional[str] = ...,
7070
keep_files: bool = ...,
@@ -91,7 +91,7 @@ def read_sql_query(
9191
categories: Optional[List[str]] = ...,
9292
chunksize: Literal[True],
9393
s3_output: Optional[str] = ...,
94-
workgroup: Optional[str] = ...,
94+
workgroup: str = ...,
9595
encryption: Optional[str] = ...,
9696
kms_key: Optional[str] = ...,
9797
keep_files: bool = ...,
@@ -118,7 +118,7 @@ def read_sql_query(
118118
categories: Optional[List[str]] = ...,
119119
chunksize: bool,
120120
s3_output: Optional[str] = ...,
121-
workgroup: Optional[str] = ...,
121+
workgroup: str = ...,
122122
encryption: Optional[str] = ...,
123123
kms_key: Optional[str] = ...,
124124
keep_files: bool = ...,
@@ -145,7 +145,7 @@ def read_sql_query(
145145
categories: Optional[List[str]] = ...,
146146
chunksize: int,
147147
s3_output: Optional[str] = ...,
148-
workgroup: Optional[str] = ...,
148+
workgroup: str = ...,
149149
encryption: Optional[str] = ...,
150150
kms_key: Optional[str] = ...,
151151
keep_files: bool = ...,
@@ -172,7 +172,7 @@ def read_sql_query(
172172
categories: Optional[List[str]] = ...,
173173
chunksize: Optional[Union[int, bool]],
174174
s3_output: Optional[str] = ...,
175-
workgroup: Optional[str] = ...,
175+
workgroup: str = ...,
176176
encryption: Optional[str] = ...,
177177
kms_key: Optional[str] = ...,
178178
keep_files: bool = ...,
@@ -199,7 +199,7 @@ def read_sql_table(
199199
categories: Optional[List[str]] = ...,
200200
chunksize: Union[None, Literal[False]] = ...,
201201
s3_output: Optional[str] = ...,
202-
workgroup: Optional[str] = ...,
202+
workgroup: str = ...,
203203
encryption: Optional[str] = ...,
204204
kms_key: Optional[str] = ...,
205205
keep_files: bool = ...,
@@ -223,7 +223,7 @@ def read_sql_table(
223223
categories: Optional[List[str]] = ...,
224224
chunksize: Literal[True],
225225
s3_output: Optional[str] = ...,
226-
workgroup: Optional[str] = ...,
226+
workgroup: str = ...,
227227
encryption: Optional[str] = ...,
228228
kms_key: Optional[str] = ...,
229229
keep_files: bool = ...,
@@ -247,7 +247,7 @@ def read_sql_table(
247247
categories: Optional[List[str]] = ...,
248248
chunksize: bool,
249249
s3_output: Optional[str] = ...,
250-
workgroup: Optional[str] = ...,
250+
workgroup: str = ...,
251251
encryption: Optional[str] = ...,
252252
kms_key: Optional[str] = ...,
253253
keep_files: bool = ...,
@@ -271,7 +271,7 @@ def read_sql_table(
271271
categories: Optional[List[str]] = ...,
272272
chunksize: int,
273273
s3_output: Optional[str] = ...,
274-
workgroup: Optional[str] = ...,
274+
workgroup: str = ...,
275275
encryption: Optional[str] = ...,
276276
kms_key: Optional[str] = ...,
277277
keep_files: bool = ...,
@@ -295,7 +295,7 @@ def read_sql_table(
295295
categories: Optional[List[str]] = ...,
296296
chunksize: Optional[Union[int, bool]],
297297
s3_output: Optional[str] = ...,
298-
workgroup: Optional[str] = ...,
298+
workgroup: str = ...,
299299
encryption: Optional[str] = ...,
300300
kms_key: Optional[str] = ...,
301301
keep_files: bool = ...,
@@ -315,7 +315,7 @@ def unload(
315315
compression: Optional[str] = ...,
316316
field_delimiter: Optional[str] = ...,
317317
partitioned_by: Optional[List[str]] = ...,
318-
workgroup: Optional[str] = ...,
318+
workgroup: str = ...,
319319
encryption: Optional[str] = ...,
320320
kms_key: Optional[str] = ...,
321321
boto3_session: Optional[boto3.Session] = ...,

‎awswrangler/athena/_statements.py‎

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def _does_statement_exist(
3535
def create_prepared_statement(
3636
sql: str,
3737
statement_name: str,
38-
workgroup: Optional[str] = None,
38+
workgroup: str = "primary",
3939
mode: Literal["update", "error"] = "update",
4040
boto3_session: Optional[boto3.Session] = None,
4141
) -> None:
@@ -50,8 +50,8 @@ def create_prepared_statement(
5050
The query string for the prepared statement.
5151
statement_name : str
5252
The name of the prepared statement.
53-
workgroup : str, optional
54-
The name of the workgroup to which the prepared statement belongs.
53+
workgroup : str
54+
The name of the workgroup to which the prepared statement belongs. Primary by default.
5555
mode: str
5656
Determines the behaviour if the prepared statement already exists:
5757
@@ -72,7 +72,6 @@ def create_prepared_statement(
7272
raise exceptions.InvalidArgumentValue("`mode` must be one of 'update' or 'error'.")
7373

7474
athena_client = _utils.client("athena", session=boto3_session)
75-
workgroup = workgroup if workgroup else "primary"
7675

7776
already_exists = _does_statement_exist(statement_name, workgroup, athena_client)
7877
if already_exists and mode == "error":
@@ -95,16 +94,14 @@ def create_prepared_statement(
9594

9695

9796
@apply_configs
98-
def list_prepared_statements(
99-
workgroup: Optional[str] = None, boto3_session: Optional[boto3.Session] = None
100-
) -> List[str]:
97+
def list_prepared_statements(workgroup: str = "primary", boto3_session: Optional[boto3.Session] = None) -> List[str]:
10198
"""
10299
List the prepared statements in the specified workgroup.
103100
104101
Parameters
105102
----------
106-
workgroup: str, optional
107-
The name of the workgroup to which the prepared statement belongs.
103+
workgroup: str
104+
The name of the workgroup to which the prepared statement belongs. Primary by default.
108105
boto3_session : boto3.Session(), optional
109106
Boto3 Session. The default boto3 session will be used if boto3_session receive None.
110107
@@ -115,7 +112,6 @@ def list_prepared_statements(
115112
Each item is a dictionary with the keys ``StatementName`` and ``LastModifiedTime``.
116113
"""
117114
athena_client = _utils.client("athena", session=boto3_session)
118-
workgroup = workgroup if workgroup else "primary"
119115

120116
response = athena_client.list_prepared_statements(WorkGroup=workgroup)
121117
statements = response["PreparedStatements"]
@@ -130,7 +126,7 @@ def list_prepared_statements(
130126
@apply_configs
131127
def delete_prepared_statement(
132128
statement_name: str,
133-
workgroup: Optional[str] = None,
129+
workgroup: str = "primary",
134130
boto3_session: Optional[boto3.Session] = None,
135131
) -> None:
136132
"""
@@ -143,7 +139,7 @@ def delete_prepared_statement(
143139
statement_name : str
144140
The name of the prepared statement.
145141
workgroup : str, optional
146-
The name of the workgroup to which the prepared statement belongs.
142+
The name of the workgroup to which the prepared statement belongs. Primary by default.
147143
boto3_session : boto3.Session(), optional
148144
Boto3 Session. The default boto3 session will be used if boto3_session receive None.
149145

‎awswrangler/athena/_utils.py‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,7 @@ def _start_query_execution(
136136
return response["QueryExecutionId"]
137137

138138

139-
def _get_workgroup_config(session: Optional[boto3.Session] = None, workgroup: Optional[str] = None) -> _WorkGroupConfig:
139+
def _get_workgroup_config(session: Optional[boto3.Session] = None, workgroup: str = "primary") -> _WorkGroupConfig:
140140
enforced: bool
141141
wg_s3_output: Optional[str]
142142
wg_encryption: Optional[str]
@@ -472,7 +472,7 @@ def repair_table(
472472
database: Optional[str] = None,
473473
data_source: Optional[str] = None,
474474
s3_output: Optional[str] = None,
475-
workgroup: Optional[str] = None,
475+
workgroup: str = "primary",
476476
encryption: Optional[str] = None,
477477
kms_key: Optional[str] = None,
478478
athena_query_wait_polling_delay: float = _QUERY_WAIT_POLLING_DELAY,
@@ -501,8 +501,8 @@ def repair_table(
501501
Data Source / Catalog name. If None, 'AwsDataCatalog' is used.
502502
s3_output : str, optional
503503
AWS S3 path.
504-
workgroup : str, optional
505-
Athena workgroup.
504+
workgroup : str
505+
Athena workgroup. Primary by default.
506506
encryption : str, optional
507507
None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'.
508508
kms_key : str, optional
@@ -552,7 +552,7 @@ def describe_table(
552552
table: str,
553553
database: Optional[str] = None,
554554
s3_output: Optional[str] = None,
555-
workgroup: Optional[str] = None,
555+
workgroup: str = "primary",
556556
encryption: Optional[str] = None,
557557
kms_key: Optional[str] = None,
558558
athena_query_wait_polling_delay: float = _QUERY_WAIT_POLLING_DELAY,
@@ -577,8 +577,8 @@ def describe_table(
577577
AWS Glue/Athena database name.
578578
s3_output : str, optional
579579
AWS S3 path.
580-
workgroup : str, optional
581-
Athena workgroup.
580+
workgroup : str
581+
Athena workgroup. Primary by default.
582582
encryption : str, optional
583583
None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'.
584584
kms_key : str, optional
@@ -641,7 +641,7 @@ def create_ctas_table( # pylint: disable=too-many-locals
641641
bucketing_info: Optional[typing.BucketingInfoTuple] = None,
642642
field_delimiter: Optional[str] = None,
643643
schema_only: bool = False,
644-
workgroup: Optional[str] = None,
644+
workgroup: str = "primary",
645645
data_source: Optional[str] = None,
646646
encryption: Optional[str] = None,
647647
kms_key: Optional[str] = None,
@@ -686,8 +686,8 @@ def create_ctas_table( # pylint: disable=too-many-locals
686686
The single-character field delimiter for files in CSV, TSV, and text files.
687687
schema_only : bool, optional
688688
_description_, by default False
689-
workgroup : Optional[str], optional
690-
Athena workgroup.
689+
workgroup : str
690+
Athena workgroup. Primary by default.
691691
data_source : Optional[str], optional
692692
Data Source / Catalog name. If None, 'AwsDataCatalog' is used.
693693
encryption : str, optional
@@ -856,7 +856,7 @@ def show_create_table(
856856
table: str,
857857
database: Optional[str] = None,
858858
s3_output: Optional[str] = None,
859-
workgroup: Optional[str] = None,
859+
workgroup: str = "primary",
860860
encryption: Optional[str] = None,
861861
kms_key: Optional[str] = None,
862862
athena_query_wait_polling_delay: float = _QUERY_WAIT_POLLING_DELAY,
@@ -880,8 +880,8 @@ def show_create_table(
880880
AWS Glue/Athena database name.
881881
s3_output : str, optional
882882
AWS S3 path.
883-
workgroup : str, optional
884-
Athena workgroup.
883+
workgroup : str
884+
Athena workgroup. Primary by default.
885885
encryption : str, optional
886886
None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'.
887887
kms_key : str, optional

‎awswrangler/athena/_write_iceberg.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,7 @@ def to_iceberg(
196196
partition_cols: Optional[List[str]] = None,
197197
keep_files: bool = True,
198198
data_source: Optional[str] = None,
199-
workgroup: Optional[str] = None,
199+
workgroup: str = "primary",
200200
encryption: Optional[str] = None,
201201
kms_key: Optional[str] = None,
202202
boto3_session: Optional[boto3.Session] = None,
@@ -237,8 +237,8 @@ def to_iceberg(
237237
Whether staging files produced by Athena are retained. 'True' by default.
238238
data_source : str, optional
239239
Data Source / Catalog name. If None, 'AwsDataCatalog' will be used by default.
240-
workgroup : str, optional
241-
Athena workgroup.
240+
workgroup : str
241+
Athena workgroup. Primary by default.
242242
encryption : str, optional
243243
Valid values: [None, 'SSE_S3', 'SSE_KMS']. Notice: 'CSE_KMS' is not supported.
244244
kms_key : str, optional

‎awswrangler/neptune/_neptune.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,8 @@ def to_property_graph(
176176
raise exceptions.InvalidArgumentValue(
177177
"DataFrame must contain at least a ~id and a ~label column to be saved to Amazon Neptune"
178178
)
179+
if df.empty:
180+
raise exceptions.EmptyDataFrame("DataFrame cannot be empty.")
179181

180182
# Loop through items in the DF
181183
for index, row in df.iterrows():

0 commit comments

Comments
 (0)