import datetime
import os
import shutil
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Type, Union
from dataclasses_json import DataClassJsonMixin
from flytekit import PythonInstanceTask, lazy_module
from flytekit.core.context_manager import FlyteContext
from flytekit.extend import Interface
from flytekit.loggers import logger
from flytekit.types.file.file import FlyteFile
from flytekit.types.schema import FlyteSchema
ge = lazy_module("great_expectations")
[docs]
@dataclass
class BatchRequestConfig(DataClassJsonMixin):
"""
Use this configuration to configure Batch Request. A BatchRequest can either be
a simple BatchRequest or a RuntimeBatchRequest.
Args:
data_connector_query: query to request data batch
runtime_parameters: parameters to be passed at runtime
batch_identifiers: identifiers to identify the data batch
batch_spec_passthrough: reader method if your file doesn't have an extension
"""
data_connector_query: Optional[Dict[str, Any]] = None
runtime_parameters: Optional[Dict[str, Any]] = None
batch_identifiers: Optional[Dict[str, str]] = None
batch_spec_passthrough: Optional[Dict[str, Any]] = None
[docs]
class GreatExpectationsTask(PythonInstanceTask[BatchRequestConfig]):
"""
This task can be used to validate your data.
You can use this when you want to validate your data within the task or workflow.
If you want to validate your data as and when the type is given, use the `GreatExpectationsType`.
Args:
name: name of the task
datasource_name: tell where your data lives and how to get it
expectation_suite_name: suite which consists of the data expectations
data_connector_name: connector to identify data batches
inputs: inputs to pass to the execute() method
data_asset_name: name of the data asset (to be used for RuntimeBatchRequest)
local_file_path: dataset file path useful for FlyteFile and FlyteSchema
checkpoint_params: optional SimpleCheckpoint parameters
task_config: batchrequest config
context_root_dir: directory in which GreatExpectations' configuration resides
TODO: Connect Data Docs to Flyte Console.
"""
_TASK_TYPE = "great_expectations"
_RUNTIME_VAR_NAME = "runtime"
def __init__(
self,
name: str,
datasource_name: str,
expectation_suite_name: str,
data_connector_name: str,
inputs: Dict[str, Type],
data_asset_name: Optional[str] = None,
outputs: Optional[Dict[str, Type]] = None,
local_file_path: Optional[str] = None,
checkpoint_params: Optional[Dict[str, Union[str, List[str]]]] = None,
task_config: BatchRequestConfig = None,
context_root_dir: str = "./great_expectations",
**kwargs,
):
self._datasource_name = datasource_name
self._data_connector_name = data_connector_name
self._expectation_suite_name = expectation_suite_name
self._batch_request_config = task_config
self._context_root_dir = context_root_dir
self._data_asset_name = data_asset_name
"""
local_file_path is a must in two scenarios:
* When using FlyteSchema
* When using FlyteFile for remote paths
This is because base directory which has the dataset file 'must' be given in GreatExpectations' config file
"""
self._local_file_path = local_file_path
self._checkpoint_params = checkpoint_params
outputs = {"result": Dict[Any, Any]}
super(GreatExpectationsTask, self).__init__(
name=name,
task_config=task_config,
task_type=self._TASK_TYPE,
interface=Interface(inputs=inputs, outputs=outputs),
**kwargs,
)
def _flyte_file(self, dataset: FlyteFile) -> str:
if not self._local_file_path:
raise ValueError("local_file_path is missing!")
shutil.copy(dataset, self._local_file_path)
return os.path.basename(dataset)
def _flyte_schema(self, dataset: FlyteSchema) -> str:
if not self._local_file_path:
raise ValueError("local_file_path is missing!")
# copy parquet file to user-given directory
FlyteContext.current_context().file_access.get_data(
dataset.remote_path, self._local_file_path, is_multipart=True
)
return os.path.basename(self._local_file_path)
[docs]
def execute(self, **kwargs) -> Any:
context = ge.data_context.DataContext(self._context_root_dir) # type: ignore
if len(self.python_interface.inputs.keys()) != 1:
raise TypeError("Expected one input argument to validate the dataset")
dataset_key = list(self.python_interface.inputs.keys())[0]
dataset = kwargs[dataset_key]
datatype = self.python_interface.inputs[dataset_key]
if not issubclass(datatype, (FlyteFile, FlyteSchema, str)):
raise TypeError("'dataset' has to have FlyteFile/FlyteSchema/str datatype")
# determine the type of data connector
selected_datasource = list(filter(lambda x: x["name"] == self._datasource_name, context.list_datasources()))
if not selected_datasource:
raise ValueError("Datasource doesn't exist!")
data_connector_class_lookup = {
data_connector_name: data_connector_class["class_name"]
for data_connector_name, data_connector_class in selected_datasource[0]["data_connectors"].items()
}
specified_data_connector_class = data_connector_class_lookup[self._data_connector_name]
is_runtime = False
if specified_data_connector_class == "RuntimeDataConnector":
is_runtime = True
if not self._data_asset_name:
raise ValueError("data_asset_name has to be given in a RuntimeBatchRequest")
# FlyteFile
if issubclass(datatype, FlyteFile):
dataset = self._flyte_file(dataset)
# FlyteSchema
# convert schema to parquet file
if issubclass(datatype, FlyteSchema) and not is_runtime:
dataset = self._flyte_schema(dataset)
# minimalistic batch request
final_batch_request = {
"data_asset_name": self._data_asset_name if is_runtime else dataset,
"datasource_name": self._datasource_name,
"data_connector_name": self._data_connector_name,
}
# Great Expectations' RuntimeBatchRequest
if self._batch_request_config and (self._batch_request_config.runtime_parameters or is_runtime):
final_batch_request.update(
{
"runtime_parameters": self._batch_request_config.runtime_parameters
if self._batch_request_config.runtime_parameters
else {},
"batch_identifiers": self._batch_request_config.batch_identifiers,
"batch_spec_passthrough": self._batch_request_config.batch_spec_passthrough,
}
)
if is_runtime and issubclass(datatype, str):
final_batch_request["runtime_parameters"]["query"] = dataset
elif is_runtime and issubclass(datatype, FlyteSchema):
# if execution engine is SparkDF, transform the data to pyspark.sql.dataframe.DataFrame, else transform the data
# to the default pandas.dataframe
if selected_datasource[0]["execution_engine"]["class_name"] == "SparkDFExecutionEngine":
import pyspark
final_batch_request["runtime_parameters"]["batch_data"] = dataset.open(
pyspark.sql.dataframe.DataFrame
).all()
else:
final_batch_request["runtime_parameters"]["batch_data"] = dataset.open().all()
else:
raise AssertionError("Can only use runtime_parameters for query(str)/schema data")
# Great Expectations' BatchRequest
elif self._batch_request_config:
final_batch_request.update(
{
"data_connector_query": self._batch_request_config.data_connector_query,
"batch_spec_passthrough": self._batch_request_config.batch_spec_passthrough,
}
)
if self._checkpoint_params:
checkpoint = ge.checkpoint.SimpleCheckpoint(
f"_tmp_checkpoint_{self._expectation_suite_name}",
context,
**self._checkpoint_params,
)
else:
checkpoint = ge.checkpoint.SimpleCheckpoint(
f"_tmp_checkpoint_{self._expectation_suite_name}",
context,
)
# identify every run uniquely
run_id = ge.core.run_identifier.RunIdentifier(
**{
"run_name": self._datasource_name + "_run",
"run_time": datetime.datetime.now(datetime.timezone.utc),
}
)
checkpoint_result = checkpoint.run(
run_id=run_id,
validations=[
{
"batch_request": final_batch_request,
"expectation_suite_name": self._expectation_suite_name,
}
],
)
final_result = ge.core.util.convert_to_json_serializable(checkpoint_result.list_validation_results())[0]
result_string = ""
if final_result["success"] is False:
for every_result in final_result["results"]:
if every_result["success"] is False:
result_string += (
every_result["expectation_config"]["kwargs"]["column"]
+ " -> "
+ every_result["expectation_config"]["expectation_type"]
+ "\n"
)
# raise a Great Expectations' exception
raise ge.exceptions.ValidationError("Validation failed!\nCOLUMN\t\tFAILED EXPECTATION\n" + result_string)
logger.info("Validation succeeded!")
return final_result