Skip to content

Commit

Permalink
Fixed typing
Browse files Browse the repository at this point in the history
  • Loading branch information
svdimchenko committed Jan 14, 2025
1 parent d7311f0 commit 587f999
Showing 1 changed file with 16 additions and 7 deletions.
23 changes: 16 additions & 7 deletions metadata-ingestion/src/datahub/ingestion/source/aws/glue.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from dataclasses import dataclass, field as dataclass_field
from functools import lru_cache
from typing import (
TYPE_CHECKING,
Any,
DefaultDict,
Dict,
Expand Down Expand Up @@ -114,6 +115,12 @@
from datahub.utilities.delta import delta_type_to_hive_type
from datahub.utilities.hive_schema_to_avro import get_schema_fields_for_hive_column

if TYPE_CHECKING:
from mypy_boto3_glue.type_defs import (
DatabasePaginatorTypeDef,
TablePaginatorTypeDef,
)

logger = logging.getLogger(__name__)

DEFAULT_PLATFORM = "glue"
Expand Down Expand Up @@ -156,8 +163,8 @@ class GlueSourceConfig(
default=None,
description="The aws account id where the target glue catalog lives. If None, datahub will ingest glue in aws caller's account.",
)
catalog_name: str = Field(
default="awsdatacatalog", description="The aws athena catalog name"
catalog_name: Optional[str] = Field(
default=None, description="The aws athena catalog name"
)
ignore_resource_links: Optional[bool] = Field(
default=False,
Expand Down Expand Up @@ -715,7 +722,7 @@ def get_datajob_wu(self, node: Dict[str, Any], job_name: str) -> MetadataWorkUni

return MetadataWorkUnit(id=f'{job_name}-{node["Id"]}', mce=mce)

def get_all_databases(self) -> Iterable[Mapping[str, Any]]:
def get_all_databases(self) -> Iterable[DatabasePaginatorTypeDef]:
logger.debug("Getting all databases")
# see https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/glue/paginator/GetDatabases.html
paginator = self.glue_client.get_paginator("get_databases")
Expand Down Expand Up @@ -743,7 +750,9 @@ def get_all_databases(self) -> Iterable[Mapping[str, Any]]:
self.report.databases.processed(database["Name"])
yield database

def get_tables_from_database(self, database: Mapping[str, Any]) -> Iterable[Dict]:
def get_tables_from_database(

Check warning on line 753 in metadata-ingestion/src/datahub/ingestion/source/aws/glue.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/ingestion/source/aws/glue.py#L753

Added line #L753 was not covered by tests
self, database: DatabasePaginatorTypeDef
) -> Iterable[TablePaginatorTypeDef]:
logger.debug(f"Getting tables from database {database['Name']}")
# see https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/glue/paginator/GetTables.html
paginator = self.glue_client.get_paginator("get_tables")
Expand All @@ -770,7 +779,7 @@ def get_tables_from_database(self, database: Mapping[str, Any]) -> Iterable[Dict

def get_all_databases_and_tables(
self,
) -> Tuple[List[Mapping[str, Any]], List[Dict]]:
) -> Tuple[List[DatabasePaginatorTypeDef], List[TablePaginatorTypeDef]]:
all_databases = [*self.get_all_databases()]
all_tables = [
tables
Expand Down Expand Up @@ -1038,7 +1047,7 @@ def gen_database_key(self, database: str) -> DatabaseKey:
)

def gen_database_containers(
self, database: Mapping[str, Any]
self, database: DatabasePaginatorTypeDef
) -> Iterable[MetadataWorkUnit]:
domain_urn = self._gen_domain_urn(database["Name"])
database_container_key = self.gen_database_key(database["Name"])
Expand Down Expand Up @@ -1113,7 +1122,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
if self.extract_transforms:
yield from self._transform_extraction()

def _gen_table_wu(self, table: Dict) -> Iterable[MetadataWorkUnit]:
def _gen_table_wu(self, table: TablePaginatorTypeDef) -> Iterable[MetadataWorkUnit]:

Check warning on line 1125 in metadata-ingestion/src/datahub/ingestion/source/aws/glue.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/ingestion/source/aws/glue.py#L1125

Added line #L1125 was not covered by tests
database_name = table["DatabaseName"]
table_name = table["Name"]
full_table_name = (

Check warning on line 1128 in metadata-ingestion/src/datahub/ingestion/source/aws/glue.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/ingestion/source/aws/glue.py#L1128

Added line #L1128 was not covered by tests
Expand Down

0 comments on commit 587f999

Please sign in to comment.