Merge pull request #525 from DagsHub/enhancement/query-time

Enhancement: Add QueryResult.query_data_time

dagshub/data_engine/client/data_client.py

@@ -133,6 +133,7 @@ def sample(self, datasource: "Datasource", n: Optional[int], include_metadata: b
  new_entries = QueryResult.from_gql_query(resp, datasource)
  res.entries += new_entries.entries
  res.fields = new_entries.fields
+ res.query_data_time = new_entries.query_data_time
  left -= take
  progress.update(total_task, advance=len(res.entries), refresh=True)
  return res
@@ -159,6 +160,7 @@ def _get_all(self, datasource: "Datasource", include_metadata: bool) -> QueryRes
  new_entries = QueryResult.from_gql_query(resp, datasource)
  res.entries += new_entries.entries
  res.fields = new_entries.fields
+ res.query_data_time = new_entries.query_data_time
  progress.update(total_task, advance=len(new_entries.entries), refresh=True)
 
  return res

dagshub/data_engine/client/gql_queries.py

@@ -72,7 +72,7 @@ def datasource_query(include_metadata: bool, introspection: "TypesIntrospection"
  )
  fields = [
  f"edges {{ node {{ id path {metadata_fields} }} }}",
- "pageInfo { hasNextPage endCursor }",
+ "pageInfo { hasNextPage endCursor } queryDataTime",
  ]
  if include_metadata:
  fields.append("selectFields { name originalName autoGenerated valueType asOf multiple tags }")

dagshub/data_engine/model/query_result.py

@@ -1,3 +1,4 @@
+import datetime
 import inspect
 import logging
 from collections import Counter, defaultdict
@@ -72,6 +73,7 @@ class QueryResult:
  _entries: List[Datapoint]
  datasource: "Datasource"
  fields: List[MetadataSelectFieldSchema]
+ query_data_time: Optional[datetime.datetime] = None
  _datapoint_path_lookup: Dict[str, Datapoint] = field(init=False)
 
  def __post_init__(self):
@@ -133,12 +135,9 @@ def from_gql_query(query_resp: Dict[str, Any], datasource: "Datasource") -> "Que
  if edges is None:
  edges = []
  datapoints = [Datapoint.from_gql_edge(edge, datasource, fields) for edge in edges]
+ query_data_time = datetime.datetime.fromtimestamp(query_resp.get("queryDataTime"), tz=datetime.timezone.utc)
 
- return QueryResult(
- datapoints,
- datasource,
- fields,
- )
+ return QueryResult(_entries=datapoints, datasource=datasource, fields=fields, query_data_time=query_data_time)
 
  def as_ml_dataset(self, flavor: str, **kwargs):
  """