datacontract-cli 0.10.0__py3-none-any.whl → 0.10.37__py3-none-any.whl

datacontract/engines/data_contract_test.py

@@ -0,0 +1,113 @@
+import atexit
+import os
+import tempfile
+import typing
+
+import requests
+from duckdb.duckdb import DuckDBPyConnection
+
+from datacontract.engines.data_contract_checks import create_checks
+
+if typing.TYPE_CHECKING:
+    from pyspark.sql import SparkSession
+
+from datacontract.engines.datacontract.check_that_datacontract_contains_valid_servers_configuration import (
+    check_that_datacontract_contains_valid_server_configuration,
+)
+from datacontract.engines.fastjsonschema.check_jsonschema import check_jsonschema
+from datacontract.engines.soda.check_soda_execute import check_soda_execute
+from datacontract.model.data_contract_specification import DataContractSpecification, Server
+from datacontract.model.exceptions import DataContractException
+from datacontract.model.run import ResultEnum, Run
+
+
+def execute_data_contract_test(
+    data_contract_specification: DataContractSpecification,
+    run: Run,
+    server_name: str = None,
+    spark: "SparkSession" = None,
+    duckdb_connection: DuckDBPyConnection = None,
+):
+    if data_contract_specification.models is None or len(data_contract_specification.models) == 0:
+        raise DataContractException(
+            type="lint",
+            name="Check that data contract contains models",
+            result=ResultEnum.warning,
+            reason="Models block is missing. Skip executing tests.",
+            engine="datacontract",
+        )
+    if (
+        server_name is None
+        and data_contract_specification.servers is not None
+        and len(data_contract_specification.servers) > 0
+    ):
+        server_name = list(data_contract_specification.servers.keys())[0]
+    server = get_server(data_contract_specification, server_name)
+    run.log_info(f"Running tests for data contract {data_contract_specification.id} with server {server_name}")
+    run.dataContractId = data_contract_specification.id
+    run.dataContractVersion = data_contract_specification.info.version
+    run.dataProductId = server.dataProductId
+    run.outputPortId = server.outputPortId
+    run.server = server_name
+
+    if server.type == "api":
+        server = process_api_response(run, server)
+
+    run.checks.extend(create_checks(data_contract_specification, server))
+
+    # TODO check server is supported type for nicer error messages
+    # TODO check server credentials are complete for nicer error messages
+    if server.format == "json" and server.type != "kafka":
+        check_jsonschema(run, data_contract_specification, server)
+    check_soda_execute(run, data_contract_specification, server, spark, duckdb_connection)
+
+
+def get_server(data_contract_specification: DataContractSpecification, server_name: str = None) -> Server | None:
+    """Get the server configuration from the data contract specification.
+
+    Args:
+        data_contract_specification: The data contract specification
+        server_name: Optional name of the server to use. If not provided, uses the first server.
+
+    Returns:
+        The selected server configuration
+    """
+
+    check_that_datacontract_contains_valid_server_configuration(data_contract_specification, server_name)
+
+    if server_name is not None:
+        server = data_contract_specification.servers.get(server_name)
+    else:
+        server_name = list(data_contract_specification.servers.keys())[0]
+        server = data_contract_specification.servers.get(server_name)
+    return server
+
+
+def process_api_response(run, server):
+    tmp_dir = tempfile.TemporaryDirectory(prefix="datacontract_cli_api_")
+    atexit.register(tmp_dir.cleanup)
+    headers = {}
+    if os.getenv("DATACONTRACT_API_HEADER_AUTHORIZATION") is not None:
+        headers["Authorization"] = os.getenv("DATACONTRACT_API_HEADER_AUTHORIZATION")
+    try:
+        response = requests.get(server.location, headers=headers)
+        response.raise_for_status()
+    except requests.exceptions.RequestException as e:
+        raise DataContractException(
+            type="connection",
+            name="API server connection error",
+            result=ResultEnum.error,
+            reason=f"Failed to fetch API response from {server.location}: {e}",
+            engine="datacontract",
+        )
+    with open(f"{tmp_dir.name}/api_response.json", "w") as f:
+        f.write(response.text)
+    run.log_info(f"Saved API response to {tmp_dir.name}/api_response.json")
+    server = Server(
+        type="local",
+        format="json",
+        path=f"{tmp_dir.name}/api_response.json",
+        dataProductId=server.dataProductId,
+        outputPortId=server.outputPortId,
+    )
+    return server

datacontract/engines/datacontract/check_that_datacontract_contains_valid_servers_configuration.py

@@ -1,12 +1,11 @@
 from datacontract.model.data_contract_specification import DataContractSpecification
 from datacontract.model.exceptions import DataContractException
-from datacontract.model.run import Run
 
 
 def check_that_datacontract_contains_valid_server_configuration(
-    run: Run, data_contract: DataContractSpecification, server_name: str
+    data_contract: DataContractSpecification, server_name: str | None
 ):
-    if data_contract.servers is None:
+    if data_contract.servers is None or len(data_contract.servers) == 0:
         raise DataContractException(
             type="lint",
             name="Check that data contract contains valid server configuration",

datacontract/engines/datacontract/check_that_datacontract_file_exists.py

@@ -1,6 +1,6 @@
 import os
 
-from datacontract.model.run import Run, Check
+from datacontract.model.run import Check, Run
 
 
 def check_that_datacontract_file_exists(run: Run, file_path: str):

datacontract/engines/fastjsonschema/check_jsonschema.py

@@ -1,32 +1,115 @@
+import glob
 import json
 import logging
 import os
+import threading
+from typing import Any, Callable, Generator, List, Optional
 
 import fastjsonschema
+from fastjsonschema import JsonSchemaValueException
 
 from datacontract.engines.fastjsonschema.s3.s3_read_files import yield_s3_files
 from datacontract.export.jsonschema_converter import to_jsonschema
-from datacontract.model.data_contract_specification import \
-    DataContractSpecification, Server
+from datacontract.model.data_contract_specification import DataContractSpecification, Server
 from datacontract.model.exceptions import DataContractException
-from datacontract.model.run import Run, Check
+from datacontract.model.run import Check, ResultEnum, Run
 
+# Thread-safe cache for primaryKey fields.
+_primary_key_cache = {}
+_cache_lock = threading.Lock()
 
-def validate_json_stream(model_name, validate, json_stream):
+
+def get_primary_key_field(schema: dict, model_name: str) -> Optional[str]:
+    # Check cache first.
+    with _cache_lock:
+        cached_value = _primary_key_cache.get(model_name)
+        if cached_value is not None:
+            return cached_value
+
+    # Find primaryKey field.
+    fields = schema.get("properties", {})
+    for field_name, attributes in fields.items():
+        if attributes.get("primaryKey", False):
+            # Cache the result before returning.
+            with _cache_lock:
+                _primary_key_cache[model_name] = field_name
+            return field_name
+
+    # Return None if no primary key was found.
+    return None
+
+
+def get_primary_key_value(schema: dict, model_name: str, json_object: dict) -> Optional[str]:
+    # Get the `primaryKey` field.
+    primary_key_field = get_primary_key_field(schema, model_name)
+    if not primary_key_field:
+        return None
+
+    # Return the value of the `primaryKey` field in the JSON object.
+    return json_object.get(primary_key_field)
+
+
+def process_exceptions(run, exceptions: List[DataContractException]):
+    if not exceptions:
+        return
+
+    # Define the maximum number of errors to process (can be adjusted by defining an ENV variable).
     try:
-        logging.info("Validating JSON")
-        for json_obj in json_stream:
+        error_limit = int(os.getenv("DATACONTRACT_MAX_ERRORS", 500))
+    except ValueError:
+        # Fallback to default if environment variable is invalid.
+        error_limit = 500
+
+    # Calculate the effective limit to avoid index out of range
+    limit = min(len(exceptions), error_limit)
+
+    # Add all exceptions up to the limit - 1 to `run.checks`.
+    DEFAULT_ERROR_MESSAGE = "An error occurred during validation phase. See the logs for more details."
+    run.checks.extend(
+        [
+            Check(
+                type=exception.type,
+                name=exception.name,
+                result=exception.result,
+                reason=exception.reason,
+                model=exception.model,
+                engine=exception.engine,
+                message=exception.message or DEFAULT_ERROR_MESSAGE,
+            )
+            for exception in exceptions[: limit - 1]
+        ]
+    )
+
+    # Raise the last exception within the limit.
+    last_exception = exceptions[limit - 1]
+    raise last_exception
+
+
+def validate_json_stream(
+    schema: dict, model_name: str, validate: Callable, json_stream: Generator[Any, Any, None]
+) -> List[DataContractException]:
+    logging.info(f"Validating JSON stream for model: '{model_name}'.")
+    exceptions: List[DataContractException] = []
+    for json_obj in json_stream:
+        try:
             validate(json_obj)
-        return True
-    except fastjsonschema.JsonSchemaValueException as e:
-        raise DataContractException(
-            type="schema",
-            name="Check that JSON has valid schema",
-            model=model_name,
-            reason=e.message,
-            engine="jsonschema",
-            original_exception=e,
-        )
+        except JsonSchemaValueException as e:
+            logging.warning(f"Validation failed for JSON object with type: '{model_name}'.")
+            primary_key_value = get_primary_key_value(schema, model_name, json_obj)
+            exceptions.append(
+                DataContractException(
+                    type="schema",
+                    name="Check that JSON has valid schema",
+                    result=ResultEnum.failed,
+                    reason=f"{f'#{primary_key_value}: ' if primary_key_value is not None else ''}{e.message}",
+                    model=model_name,
+                    engine="jsonschema",
+                    message=e.message,
+                )
+            )
+    if not exceptions:
+        logging.info(f"All JSON objects in the stream passed validation for model: '{model_name}'.")
+    return exceptions
 
 
 def read_json_lines(file):
@@ -60,42 +143,64 @@ def read_json_file_content(file_content: str):
     yield json.loads(file_content)
 
 
-def process_json_file(run, model_name, validate, file, delimiter):
+def process_json_file(run, schema, model_name, validate, file, delimiter):
     if delimiter == "new_line":
         json_stream = read_json_lines(file)
     elif delimiter == "array":
         json_stream = read_json_array(file)
     else:
         json_stream = read_json_file(file)
-    validate_json_stream(model_name, validate, json_stream)
 
+    # Validate the JSON stream and collect exceptions.
+    exceptions = validate_json_stream(schema, model_name, validate, json_stream)
+
+    # Handle all errors from schema validation.
+    process_exceptions(run, exceptions)
 
-def process_local_file(run, server, model_name, validate):
+
+def process_local_file(run, server, schema, model_name, validate):
     path = server.path
+    if not path:
+        raise DataContractException(
+            type="schema",
+            name="Check that JSON has valid schema",
+            result=ResultEnum.warning,
+            reason="For server with type 'local', a 'path' must be defined.",
+            engine="datacontract",
+        )
     if "{model}" in path:
         path = path.format(model=model_name)
 
+    all_files = []
     if os.path.isdir(path):
-        return process_directory(run, path, server, model_name, validate)
+        # Fetch all JSONs in the directory
+        for root, _, files in os.walk(path):
+            for file in files:
+                if file.endswith(".json"):
+                    all_files.append(os.path.join(root, file))
     else:
-        logging.info(f"Processing file {path}")
-        with open(path, "r") as file:
-            process_json_file(run, model_name, validate, file, server.delimiter)
+        # Use glob to fetch all JSONs
+        for file_path in glob.glob(path, recursive=True):
+            if os.path.isfile(file_path):
+                if file_path.endswith(".json"):
+                    all_files.append(file_path)
 
+    if not all_files:
+        raise DataContractException(
+            type="schema",
+            name="Check that JSON has valid schema",
+            result=ResultEnum.warning,
+            reason=f"No files found in '{path}'.",
+            engine="datacontract",
+        )
 
-def process_directory(run, path, server, model_name, validate):
-    success = True
-    for filename in os.listdir(path):
-        if filename.endswith(".json"):  # or make this a parameter
-            file_path = os.path.join(path, filename)
-            with open(file_path, "r") as file:
-                if not process_json_file(run, model_name, validate, file, server.delimiter):
-                    success = False
-                    break
-    return success
+    for file in all_files:
+        logging.info(f"Processing file: {file}")
+        with open(file, "r") as f:
+            process_json_file(run, schema, model_name, validate, f, server.delimiter)
 
 
-def process_s3_file(server, model_name, validate):
+def process_s3_file(run, server, schema, model_name, validate):
     s3_endpoint_url = server.endpointUrl
     s3_location = server.location
     if "{model}" in s3_location:
@@ -114,12 +219,16 @@ def process_s3_file(server, model_name, validate):
         raise DataContractException(
             type="schema",
             name="Check that JSON has valid schema",
-            result="warning",
+            result=ResultEnum.warning,
             reason=f"Cannot find any file in {s3_location}",
             engine="datacontract",
         )
 
-    return validate_json_stream(model_name, validate, json_stream)
+    # Validate the JSON stream and collect exceptions.
+    exceptions = validate_json_stream(schema, model_name, validate, json_stream)
+
+    # Handle all errors from schema validation.
+    process_exceptions(run, exceptions)
 
 
 def check_jsonschema(run: Run, data_contract: DataContractSpecification, server: Server):
@@ -131,7 +240,7 @@ def check_jsonschema(run: Run, data_contract: DataContractSpecification, server:
             Check(
                 type="schema",
                 name="Check that JSON has valid schema",
-                result="warning",
+                result=ResultEnum.warning,
                 reason="Server format is not 'json'. Skip validating jsonschema.",
                 engine="jsonschema",
             )
@@ -149,20 +258,45 @@ def check_jsonschema(run: Run, data_contract: DataContractSpecification, server:
         schema = to_jsonschema(model_name, model)
         run.log_info(f"jsonschema: {schema}")
 
-        validate = fastjsonschema.compile(schema)
+        validate = fastjsonschema.compile(
+            schema,
+            formats={"uuid": r"^[0-9a-fA-F]{8}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{12}$"},
+        )
 
         # Process files based on server type
         if server.type == "local":
-            process_local_file(run, server, model_name, validate)
+            process_local_file(run, server, schema, model_name, validate)
         elif server.type == "s3":
-            process_s3_file(server, model_name, validate)
+            process_s3_file(run, server, schema, model_name, validate)
+        elif server.type == "gcs":
+            run.checks.append(
+                Check(
+                    type="schema",
+                    name="Check that JSON has valid schema",
+                    model=model_name,
+                    result=ResultEnum.info,
+                    reason="JSON Schema check skipped for GCS, as GCS is currently not supported",
+                    engine="jsonschema",
+                )
+            )
+        elif server.type == "azure":
+            run.checks.append(
+                Check(
+                    type="schema",
+                    name="Check that JSON has valid schema",
+                    model=model_name,
+                    result=ResultEnum.info,
+                    reason="JSON Schema check skipped for azure, as azure is currently not supported",
+                    engine="jsonschema",
+                )
+            )
         else:
             run.checks.append(
                 Check(
                     type="schema",
                     name="Check that JSON has valid schema",
                     model=model_name,
-                    result="warn",
+                    result=ResultEnum.warning,
                     reason=f"Server type {server.type} not supported",
                     engine="jsonschema",
                 )
@@ -174,7 +308,7 @@ def check_jsonschema(run: Run, data_contract: DataContractSpecification, server:
                 type="schema",
                 name="Check that JSON has valid schema",
                 model=model_name,
-                result="passed",
+                result=ResultEnum.passed,
                 reason="All JSON entries are valid.",
                 engine="jsonschema",
             )

datacontract/engines/fastjsonschema/s3/s3_read_files.py

@@ -1,7 +1,8 @@
 import logging
 import os
 
-import s3fs
+from datacontract.model.exceptions import DataContractException
+from datacontract.model.run import ResultEnum
 
 
 def yield_s3_files(s3_endpoint_url, s3_location):
@@ -14,11 +15,25 @@ def yield_s3_files(s3_endpoint_url, s3_location):
 
 
 def s3_fs(s3_endpoint_url):
+    try:
+        import s3fs
+    except ImportError as e:
+        raise DataContractException(
+            type="schema",
+            result=ResultEnum.failed,
+            name="s3 extra missing",
+            reason="Install the extra s3 to use s3",
+            engine="datacontract",
+            original_exception=e,
+        )
+
     aws_access_key_id = os.getenv("DATACONTRACT_S3_ACCESS_KEY_ID")
     aws_secret_access_key = os.getenv("DATACONTRACT_S3_SECRET_ACCESS_KEY")
+    aws_session_token = os.getenv("DATACONTRACT_S3_SESSION_TOKEN")
     return s3fs.S3FileSystem(
         key=aws_access_key_id,
         secret=aws_secret_access_key,
+        token=aws_session_token,
         anon=aws_access_key_id is None,
         client_kwargs={"endpoint_url": s3_endpoint_url},
     )