diff --git a/ingestion/src/metadata/ingestion/source/database/athena/models.py b/ingestion/src/metadata/ingestion/source/database/athena/models.py index 2d9b6c962d0..d833893846d 100644 --- a/ingestion/src/metadata/ingestion/source/database/athena/models.py +++ b/ingestion/src/metadata/ingestion/source/database/athena/models.py @@ -19,59 +19,23 @@ from pydantic import BaseModel class QueryExecutionIdsResponse(BaseModel): - QueryExecutionIds: List[str] - - -class ResultReuseByAgeConfiguration(BaseModel): - Enabled: bool - - -class ResultConfiguration(BaseModel): - OutputLocation: str - - -class QueryExecutionContext(BaseModel): - Database: str - Catalog: str + QueryExecutionIds: Optional[List[str]] class Status(BaseModel): - State: str - SubmissionDateTime: datetime - CompletionDateTime: datetime - - -class ResultReuseInformation(BaseModel): - ReusedPreviousResult: bool + State: Optional[str] + SubmissionDateTime: Optional[datetime] class Statistics(BaseModel): - EngineExecutionTimeInMillis: int - DataScannedInBytes: int - TotalExecutionTimeInMillis: int - QueryQueueTimeInMillis: int - ServiceProcessingTimeInMillis: int - ResultReuseInformation: ResultReuseInformation - - -class EngineVersion(BaseModel): - SelectedEngineVersion: str - EffectiveEngineVersion: str + TotalExecutionTimeInMillis: Optional[int] class AthenaQueryExecution(BaseModel): - QueryExecutionId: str - Query: str - StatementType: str - ResultConfiguration: ResultConfiguration - ResultReuseConfiguration: dict - QueryExecutionContext: QueryExecutionContext - Status: Status - Statistics: Statistics - WorkGroup: str - EngineVersion: EngineVersion - SubstatementType: Optional[str] + Query: Optional[str] + Statistics: Optional[Statistics] + Status: Optional[Status] class AthenaQueryExecutionList(BaseModel): - QueryExecutions: List[AthenaQueryExecution] + QueryExecutions: Optional[List[AthenaQueryExecution]] diff --git a/ingestion/src/metadata/ingestion/source/database/athena/usage.py b/ingestion/src/metadata/ingestion/source/database/athena/usage.py index c3a07162d57..bfeb9401b41 100644 --- a/ingestion/src/metadata/ingestion/source/database/athena/usage.py +++ b/ingestion/src/metadata/ingestion/source/database/athena/usage.py @@ -48,11 +48,15 @@ class AthenaUsageSource(AthenaQueryParserSource, UsageSource): ), analysisDate=query.Status.SubmissionDateTime, serviceName=self.config.serviceName, - duration=query.Statistics.TotalExecutionTimeInMillis, + duration=query.Statistics.TotalExecutionTimeInMillis + if query.Statistics + else None, aborted=query.Status.State == QUERY_ABORTED_STATE, ) for query in query_list.QueryExecutions - if query.Status.SubmissionDateTime.date() >= self.start.date() + if query.Status + and query.Query + and query.Status.SubmissionDateTime.date() >= self.start.date() and self.is_not_dbt_or_om_query(query.Query) ] yield TableQueries(queries=queries)