Skip to content

Commit

Permalink
move boto3 stubs back as a build dependency
Browse files Browse the repository at this point in the history
  • Loading branch information
mikealfare committed Oct 8, 2024
1 parent 4df1c0b commit 6033045
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 23 deletions.
22 changes: 9 additions & 13 deletions dbt-athena/src/dbt/adapters/athena/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -235,7 +235,7 @@ def _get_work_group(self, work_group: str) -> GetWorkGroupOutputTypeDef:
config=get_boto3_config(num_retries=creds.effective_num_retries),
)

return athena_client.get_work_group(WorkGroup=work_group) # type:ignore
return athena_client.get_work_group(WorkGroup=work_group)

@available
def is_work_group_output_location_enforced(self) -> bool:
Expand Down Expand Up @@ -355,7 +355,7 @@ def get_glue_table(self, relation: AthenaRelation) -> Optional[GetTableResponseT
LOGGER.debug(f"Table {relation.render()} does not exists - Ignoring")
return None
raise e
return table # type:ignore
return table

@available
def get_glue_table_type(self, relation: AthenaRelation) -> Optional[TableType]:
Expand Down Expand Up @@ -581,7 +581,7 @@ def _get_one_table_for_catalog(
for idx, col in enumerate(
table["StorageDescriptor"]["Columns"] + table.get("PartitionKeys", [])
)
if self._is_current_column(col) # type:ignore
if self._is_current_column(col)
]

@staticmethod
Expand All @@ -606,9 +606,7 @@ def _get_one_table_for_non_glue_catalog(
},
}
# TODO: review this code part as TableTypeDef class does not contain "Columns" attribute
for idx, col in enumerate(
table["Columns"] + table.get("PartitionKeys", []) # type:ignore
)
for idx, col in enumerate(table["Columns"] + table.get("PartitionKeys", []))
]

def _get_one_catalog(
Expand Down Expand Up @@ -714,7 +712,7 @@ def _get_data_catalog(self, database: str) -> Optional[DataCatalogTypeDef]:
region_name=client.region_name,
config=get_boto3_config(num_retries=creds.effective_num_retries),
)
return athena.get_data_catalog(Name=database)["DataCatalog"] # type:ignore
return athena.get_data_catalog(Name=database)["DataCatalog"]
return None

@available
Expand Down Expand Up @@ -1040,7 +1038,7 @@ def persist_docs_to_glue(
need_to_update_table = True
# Save dbt description
table_input["Description"] = clean_table_description
table_parameters["comment"] = clean_table_description # type:ignore
table_parameters["comment"] = clean_table_description

# Get dbt model meta if available
meta: Dict[str, Any] = model.get("config", {}).get("meta", {})
Expand All @@ -1060,7 +1058,7 @@ def persist_docs_to_glue(
if current_meta_value is None or current_meta_value != meta_value:
need_to_update_table = True
# Save Glue table parameter
table_parameters[meta_key] = meta_value # type:ignore
table_parameters[meta_key] = meta_value
else:
LOGGER.warning(
f"Meta value for key '{meta_key}' is not supported and will be ignored"
Expand Down Expand Up @@ -1106,7 +1104,7 @@ def persist_docs_to_glue(
):
need_to_update_table = True
# Save Glue column parameter
col_obj["Parameters"][meta_key] = meta_value # type:ignore
col_obj["Parameters"][meta_key] = meta_value
else:
LOGGER.warning(
f"Column meta value for key '{meta_key}' is not supported and will be ignored"
Expand Down Expand Up @@ -1363,9 +1361,7 @@ def _get_table_input(table: TableTypeDef) -> TableInputTypeDef:
This is needed because update_table() does not accept some read-only fields of table dictionary
returned by get_table() method.
"""
return {
k: v for k, v in table.items() if k in TableInputTypeDef.__annotations__
} # type:ignore
return {k: v for k, v in table.items() if k in TableInputTypeDef.__annotations__}

@available
def run_query_with_partitions_limit_catching(self, sql: str) -> str:
Expand Down
18 changes: 8 additions & 10 deletions dbt-athena/src/dbt/adapters/athena/lakeformation.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,18 +49,16 @@ def process_lf_tags_database(self) -> None:
if self.lf_tags:
database_resource = {"Database": {"Name": self.database}}
response = self.lf_client.add_lf_tags_to_resource(
Resource=database_resource, # type:ignore
Resource=database_resource,
LFTags=[{"TagKey": k, "TagValues": [v]} for k, v in self.lf_tags.items()],
)
self._parse_and_log_lf_response(response, None, self.lf_tags)

def process_lf_tags(self) -> None:
table_resource = {"Table": {"DatabaseName": self.database, "Name": self.table}}
existing_lf_tags = self.lf_client.get_resource_lf_tags(
Resource=table_resource # type:ignore
)
existing_lf_tags = self.lf_client.get_resource_lf_tags(Resource=table_resource)
self._remove_lf_tags_columns(existing_lf_tags)
self._apply_lf_tags_table(table_resource, existing_lf_tags) # type:ignore
self._apply_lf_tags_table(table_resource, existing_lf_tags)
self._apply_lf_tags_columns()

@staticmethod
Expand Down Expand Up @@ -103,7 +101,7 @@ def _remove_lf_tags_columns(self, existing_lf_tags: GetResourceLFTagsResponseTyp
}
}
response = self.lf_client.remove_lf_tags_from_resource(
Resource=resource, # type:ignore
Resource=resource,
LFTags=[{"TagKey": tag_key, "TagValues": [tag_value]}],
)
self._parse_and_log_lf_response(
Expand Down Expand Up @@ -131,7 +129,7 @@ def _apply_lf_tags_table(
logger.debug(f"CONFIG TAGS: {self.lf_tags}")

to_remove = LfTagsManager._table_tags_to_remove(
lf_tags_table, self.lf_tags, self.lf_inherited_tags # type:ignore
lf_tags_table, self.lf_tags, self.lf_inherited_tags
)

logger.debug(f"TAGS TO REMOVE: {to_remove}")
Expand Down Expand Up @@ -161,7 +159,7 @@ def _apply_lf_tags_columns(self) -> None:
}
}
response = self.lf_client.add_lf_tags_to_resource(
Resource=resource, # type:ignore
Resource=resource,
LFTags=[{"TagKey": tag_key, "TagValues": [tag_value]}],
)
self._parse_and_log_lf_response(response, columns, {tag_key: tag_value})
Expand Down Expand Up @@ -238,8 +236,8 @@ def get_filters(self) -> Dict[str, DataCellsFilterTypeDef]:
"Name": self.table,
}
return {
f["Name"]: f # type:ignore
for f in self.lf_client.list_data_cells_filter(Table=table_resource)[ # type:ignore
f["Name"]: f
for f in self.lf_client.list_data_cells_filter(Table=table_resource)[
"DataCellsFilters"
]
}
Expand Down

0 comments on commit 6033045

Please sign in to comment.