Fix #11352: Fix athena usage models (#11378)

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]]

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)