openmetadata-ingestion 1.5.13.1__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

metadata/ingestion/source/database/database_service.py

@@ -13,7 +13,7 @@ Base class for ingesting database services
 """
 import traceback
 from abc import ABC, abstractmethod
-from typing import Any, Iterable, List, Optional, Set, Tuple, Union
+from typing import Any, Iterable, List, Optional, Set, Tuple
 
 from pydantic import BaseModel, Field
 from sqlalchemy.engine import Inspector
@@ -23,7 +23,6 @@ from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequ
 from metadata.generated.schema.api.data.createDatabaseSchema import (
     CreateDatabaseSchemaRequest,
 )
-from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
 from metadata.generated.schema.api.data.createStoredProcedure import (
     CreateStoredProcedureRequest,
 )
@@ -112,13 +111,9 @@ class DatabaseServiceTopology(ServiceTopology):
             ),
         ],
         children=["database"],
-        # Note how we have `yield_view_lineage` and `yield_stored_procedure_lineage`
-        # as post_processed. This is because we cannot ensure proper lineage processing
-        # until we have finished ingesting all the metadata from the source.
         post_process=[
-            "yield_view_lineage",
-            "yield_procedure_lineage_and_queries",
             "yield_external_table_lineage",
+            "yield_table_constraints",
         ],
     )
     database: Annotated[
@@ -345,15 +340,14 @@ class DatabaseServiceSource(
         if self.source_config.includeTags:
             yield from self.yield_database_tag(database_name) or []
 
-    @abstractmethod
-    def yield_view_lineage(self) -> Iterable[Either[AddLineageRequest]]:
-        """
-        From topology.
-        Parses view definition to get lineage information
-        """
-
     def update_table_constraints(
-        self, table_constraints: List[TableConstraint], foreign_columns: []
+        self,
+        table_name,
+        schema_name,
+        db_name,
+        table_constraints: List[TableConstraint],
+        foreign_columns: [],
+        columns,
     ) -> List[TableConstraint]:
         """
         process the table constraints of all tables
@@ -381,12 +375,6 @@ class DatabaseServiceSource(
     ) -> Iterable[Either[CreateStoredProcedureRequest]]:
         """Process the stored procedure information"""
 
-    @abstractmethod
-    def yield_procedure_lineage_and_queries(
-        self,
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
-        """Extracts the lineage information from Stored Procedures"""
-
     def get_raw_database_schema_names(self) -> Iterable[str]:
         """
         fetch all schema names without any filtering.
@@ -539,7 +527,7 @@ class DatabaseServiceSource(
                 self.inspector, "get_table_owner"
             ):
                 owner_name = self.inspector.get_table_owner(
-                    connection=self.connection,  # pylint: disable=no-member.fetchall()
+                    connection=self.connection,  # pylint: disable=no-member
                     table_name=table_name,
                     schema=self.context.get().database_schema,
                 )
@@ -610,6 +598,11 @@ class DatabaseServiceSource(
         Process external table lineage
         """
 
+    def yield_table_constraints(self) -> Iterable[Either[AddLineageRequest]]:
+        """
+        Process remaining table constraints by patching the table
+        """
+
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)
         result = test_connection_fn(

metadata/ingestion/source/database/databricks/metadata.py

@@ -112,6 +112,12 @@ _type_map.update(
     }
 )
 
+
+def format_schema_name(schema):
+    # Adds back quotes(``) if hyphen(-) in schema name
+    return f"`{schema}`" if "-" in schema else schema
+
+
 # This method is from hive dialect originally but
 # is overridden to optimize DESCRIBE query execution
 def _get_table_columns(self, connection, table_name, schema, db_name):
@@ -152,6 +158,7 @@ def _get_table_columns(self, connection, table_name, schema, db_name):
 
 def _get_column_rows(self, connection, table_name, schema, db_name):
     # get columns and strip whitespace
+    schema = format_schema_name(schema=schema)
     table_columns = _get_table_columns(  # pylint: disable=protected-access
         self, connection, table_name, schema, db_name
     )
@@ -386,6 +393,7 @@ def get_table_type(self, connection, database, schema, table):
                 database_name=database, schema_name=schema, table_name=table
             )
         else:
+            schema = format_schema_name(schema=schema)
             query = f"DESCRIBE TABLE EXTENDED {schema}.{table}"
         rows = get_table_comment_result(
             self,
@@ -753,6 +761,7 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
     ) -> str:
         description = None
         try:
+            schema_name = format_schema_name(schema=schema_name)
             query = DATABRICKS_GET_TABLE_COMMENTS.format(
                 database_name=self.context.get().database,
                 schema_name=schema_name,
@@ -807,7 +816,9 @@ class DatabricksSource(ExternalTableLineageMixin, CommonDbSourceService, MultiDB
         try:
             query = DATABRICKS_GET_TABLE_COMMENTS.format(
                 database_name=self.context.get().database,
-                schema_name=self.context.get().database_schema,
+                schema_name=format_schema_name(
+                    schema=self.context.get().database_schema
+                ),
                 table_name=table_name,
             )
             result = self.inspector.dialect.get_table_comment_result(

metadata/ingestion/source/database/databricks/service_spec.py

@@ -0,0 +1,20 @@
+from metadata.data_quality.interface.sqlalchemy.databricks.test_suite_interface import (
+    DatabricksTestSuiteInterface,
+)
+from metadata.ingestion.source.database.databricks.lineage import (
+    DatabricksLineageSource,
+)
+from metadata.ingestion.source.database.databricks.metadata import DatabricksSource
+from metadata.ingestion.source.database.databricks.usage import DatabricksUsageSource
+from metadata.profiler.interface.sqlalchemy.databricks.profiler_interface import (
+    DatabricksProfilerInterface,
+)
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=DatabricksSource,
+    lineage_source_class=DatabricksLineageSource,
+    usage_source_class=DatabricksUsageSource,
+    profiler_class=DatabricksProfilerInterface,
+    test_suite_class=DatabricksTestSuiteInterface,
+)

metadata/ingestion/source/database/datalake/clients/azure_blob.py

@@ -26,9 +26,6 @@ from metadata.utils.constants import DEFAULT_DATABASE
 
 
 class DatalakeAzureBlobClient(DatalakeBaseClient):
-    def __init__(self, client: BlobServiceClient):
-        self._client = client
-
     @classmethod
     def from_config(cls, config: AzureConfig) -> "DatalakeAzureBlobClient":
         try:

metadata/ingestion/source/database/datalake/clients/base.py

@@ -13,10 +13,19 @@
 Datalake Base Client
 """
 from abc import ABC, abstractmethod
-from typing import Callable, Iterable, Optional
+from typing import Any, Callable, Iterable, Optional
 
 
 class DatalakeBaseClient(ABC):
+    """Base DL client implementation"""
+
+    def __init__(self, client: Any, **kwargs):
+        self._client = client
+
+    @property
+    def client(self) -> Any:
+        return self._client
+
     @classmethod
     @abstractmethod
     def from_config(cls, config) -> "DatalakeBaseClient":

metadata/ingestion/source/database/datalake/clients/gcs.py

@@ -33,9 +33,11 @@ from metadata.utils.credentials import GOOGLE_CREDENTIALS, set_google_credential
 
 class DatalakeGcsClient(DatalakeBaseClient):
     def __init__(
-        self, client: storage.Client, temp_credentials_file_path_list: List[str]
+        self,
+        client: storage.Client,
+        temp_credentials_file_path_list: List[str],
     ):
-        self._client = client
+        super().__init__(client=client)
         self._temp_credentials_file_path_list = temp_credentials_file_path_list
 
     @property
@@ -49,8 +51,10 @@ class DatalakeGcsClient(DatalakeBaseClient):
         if hasattr(config.securityConfig, "gcpConfig") and isinstance(
             config.securityConfig.gcpConfig.projectId, MultipleProjectId
         ):
-            gcs_config.securityConfig.gcpConfig.projectId = SingleProjectId.parse_obj(
-                gcs_config.securityConfig.gcpConfig.projectId.root[0]
+            gcs_config.securityConfig.gcpConfig.projectId = (
+                SingleProjectId.model_validate(
+                    gcs_config.securityConfig.gcpConfig.projectId.root[0]
+                )
             )
 
         if not gcs_config.securityConfig:
@@ -89,8 +93,8 @@ class DatalakeGcsClient(DatalakeBaseClient):
         gcs_config = deepcopy(config)
 
         if hasattr(gcs_config.securityConfig, "gcpConfig"):
-            gcs_config.securityConfig.gcpConfig.projectId = SingleProjectId.parse_obj(
-                database_name
+            gcs_config.securityConfig.gcpConfig.projectId = (
+                SingleProjectId.model_validate(database_name)
             )
 
         self._client = self.get_gcs_client(gcs_config)

metadata/ingestion/source/database/datalake/clients/s3.py

@@ -25,9 +25,6 @@ from metadata.utils.s3_utils import list_s3_objects
 
 
 class DatalakeS3Client(DatalakeBaseClient):
-    def __init__(self, client):
-        self._client = client
-
     @classmethod
     def from_config(cls, config: S3Config) -> "DatalakeS3Client":
         if not config.securityConfig:

metadata/ingestion/source/database/datalake/metadata.py

@@ -14,18 +14,16 @@ DataLake connector to fetch metadata from a files stored s3, gcs and Hdfs
 """
 import json
 import traceback
-from typing import Any, Iterable, Optional, Tuple, Union
+from typing import Any, Iterable, Optional, Tuple
 
 from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequest
 from metadata.generated.schema.api.data.createDatabaseSchema import (
     CreateDatabaseSchemaRequest,
 )
-from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
 from metadata.generated.schema.api.data.createStoredProcedure import (
     CreateStoredProcedureRequest,
 )
 from metadata.generated.schema.api.data.createTable import CreateTableRequest
-from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.table import Table, TableType
@@ -240,7 +238,7 @@ class DatalakeSource(DatabaseServiceSource):
                 verbose=False,
             )
             content = json.loads(metadata_config_response)
-            metadata_entry = StorageContainerConfig.parse_obj(content)
+            metadata_entry = StorageContainerConfig.model_validate(content)
         except ReadException:
             metadata_entry = None
         if self.source_config.includeTables:
@@ -319,9 +317,6 @@ class DatalakeSource(DatabaseServiceSource):
                 )
             )
 
-    def yield_view_lineage(self) -> Iterable[Either[AddLineageRequest]]:
-        yield from []
-
     def yield_tag(
         self, schema_name: str
     ) -> Iterable[Either[OMetaTagAndClassification]]:
@@ -338,12 +333,6 @@ class DatalakeSource(DatabaseServiceSource):
     def get_stored_procedure_queries(self) -> Iterable[QueryByProcedure]:
         """Not Implemented"""
 
-    def yield_procedure_lineage_and_queries(
-        self,
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
-        """Not Implemented"""
-        yield from []
-
     def standardize_table_name(
         self, schema: str, table: str  # pylint: disable=unused-argument
     ) -> str:

metadata/ingestion/source/database/datalake/service_spec.py

@@ -0,0 +1,16 @@
+from metadata.data_quality.interface.pandas.pandas_test_suite_interface import (
+    PandasTestSuiteInterface,
+)
+from metadata.ingestion.source.database.datalake.metadata import DatalakeSource
+from metadata.profiler.interface.pandas.profiler_interface import (
+    PandasProfilerInterface,
+)
+from metadata.sampler.pandas.sampler import DatalakeSampler
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=DatalakeSource,
+    profiler_class=PandasProfilerInterface,
+    test_suite_class=PandasTestSuiteInterface,
+    sampler_class=DatalakeSampler,
+)

metadata/ingestion/source/database/db2/connection.py

@@ -12,6 +12,7 @@
 """
 Source connection handler
 """
+from pathlib import Path
 from typing import Optional
 
 from sqlalchemy.engine import Engine
@@ -32,13 +33,25 @@ from metadata.ingestion.connections.builders import (
 )
 from metadata.ingestion.connections.test_connections import test_connection_db_common
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.utils.constants import THREE_MIN
+from metadata.utils.constants import THREE_MIN, UTF_8
 
 
 def get_connection(connection: Db2Connection) -> Engine:
     """
     Create connection
     """
+    # prepare license
+    # pylint: disable=import-outside-toplevel
+    if connection.license and connection.licenseFileName:
+        import clidriver
+
+        with open(
+            Path(clidriver.__path__[0], "license", connection.licenseFileName),
+            "w",
+            encoding=UTF_8,
+        ) as file:
+            file.write(connection.license)
+
     return create_generic_db_connection(
         connection=connection,
         get_connection_url_fn=get_connection_url_common,

metadata/ingestion/source/database/db2/lineage.py

@@ -0,0 +1,46 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Db2 lineage module
+"""
+from typing import Optional
+
+from metadata.generated.schema.entity.services.connections.database.db2Connection import (
+    Db2Connection,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.lineage_source import LineageSource
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class Db2LineageSource(LineageSource):
+    """
+    Db2 lineage source implements view lineage
+    """
+
+    @classmethod
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
+        """Create class instance"""
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: Db2Connection = config.serviceConnection.root.config
+        if not isinstance(connection, Db2Connection):
+            raise InvalidSourceException(
+                f"Expected Db2Connection, but got {connection}"
+            )
+        return cls(config, metadata)

metadata/ingestion/source/database/db2/service_spec.py

@@ -0,0 +1,9 @@
+from metadata.ingestion.source.database.db2.metadata import Db2Source
+from metadata.profiler.interface.sqlalchemy.db2.profiler_interface import (
+    DB2ProfilerInterface,
+)
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=Db2Source, profiler_class=DB2ProfilerInterface
+)

metadata/ingestion/source/database/dbt/constants.py

@@ -22,6 +22,14 @@ REQUIRED_MANIFEST_KEYS = ["name", "schema", "resource_type"]
 # Based on https://schemas.getdbt.com/dbt/catalog/v1.json
 REQUIRED_CATALOG_KEYS = ["name", "type", "index"]
 
+REQUIRED_CONSTRAINT_KEYS = [
+    "type",
+    "name",
+    "expression",
+    "warn_unenforced",
+    "warn_unsupported",
+]
+
 REQUIRED_RESULTS_KEYS = {
     "status",
     "timing",
@@ -82,6 +90,7 @@ NONE_KEYWORDS_LIST = ["none", "null"]
 DBT_CATALOG_FILE_NAME = "catalog.json"
 DBT_MANIFEST_FILE_NAME = "manifest.json"
 DBT_RUN_RESULTS_FILE_NAME = "run_results"
+DBT_SOURCES_FILE_NAME = "sources.json"
 
 
 class SkipResourceTypeEnum(Enum):
@@ -91,6 +100,7 @@ class SkipResourceTypeEnum(Enum):
 
     ANALYSIS = "analysis"
     TEST = "test"
+    SOURCE = "source"
 
 
 class CompiledQueriesEnum(Enum):
@@ -127,6 +137,7 @@ class DbtTestFailureEnum(Enum):
 
     FAILURE = "failure"
     FAIL = "fail"
+    ERROR = "error"
 
 
 class DbtCommonEnum(Enum):
@@ -137,6 +148,7 @@ class DbtCommonEnum(Enum):
     OWNER = "owner"
     NODES = "nodes"
     SOURCES = "sources"
+    SOURCES_FILE = "sources_file"
     SOURCE = "source"
     RESOURCETYPE = "resource_type"
     MANIFEST_NODE = "manifest_node"

metadata/ingestion/source/database/dbt/dbt_config.py

@@ -43,6 +43,7 @@ from metadata.ingestion.source.database.dbt.constants import (
     DBT_CATALOG_FILE_NAME,
     DBT_MANIFEST_FILE_NAME,
     DBT_RUN_RESULTS_FILE_NAME,
+    DBT_SOURCES_FILE_NAME,
 )
 from metadata.ingestion.source.database.dbt.models import DbtFiles
 from metadata.readers.file.config_source_factory import get_reader
@@ -85,6 +86,7 @@ def _(config: DbtLocalConfig):
             config.dbtManifestFilePath,
             config.dbtCatalogFilePath,
             config.dbtRunResultsFilePath,
+            config.dbtSourcesFilePath,
         ]
         yield from download_dbt_files(
             blob_grouped_by_directory=blob_grouped_by_directory,
@@ -123,12 +125,22 @@ def _(config: DbtHttpConfig):
             dbt_catalog = requests.get(  # pylint: disable=missing-timeout
                 config.dbtCatalogHttpPath
             )
+
+        dbt_sources = None
+        if config.dbtSourcesHttpPath:
+            logger.debug(
+                f"Requesting [dbtSourcesHttpPath] to: {config.dbtSourcesHttpPath}"
+            )
+            dbt_sources = requests.get(  # pylint: disable=missing-timeout
+                config.dbtSourcesHttpPath
+            )
         if not dbt_manifest:
             raise DBTConfigException("Manifest file not found in file server")
         yield DbtFiles(
             dbt_catalog=dbt_catalog.json() if dbt_catalog else None,
             dbt_manifest=dbt_manifest.json(),
             dbt_run_results=[dbt_run_results.json()] if dbt_run_results else None,
+            dbt_sources=dbt_sources.json() if dbt_sources else None,
         )
     except DBTConfigException as exc:
         raise exc
@@ -164,7 +176,11 @@ def _(config: DbtCloudConfig):  # pylint: disable=too-many-locals
         logger.debug(
             "Requesting [dbt_catalog], [dbt_manifest] and [dbt_run_results] data"
         )
-        params_data = {"order_by": "-finished_at", "limit": "1", "status": "10"}
+        params_data = {
+            "order_by": "-finished_at",
+            "limit": "1",
+            "status__in": "[10,20]",
+        }
         if project_id:
             params_data["project_id"] = project_id
 
@@ -182,7 +198,7 @@ def _(config: DbtCloudConfig):  # pylint: disable=too-many-locals
             last_run = runs_data[0]
             run_id = last_run["id"]
             logger.info(
-                f"Retrieved last successful run [{str(run_id)}]: "
+                f"Retrieved last completed run [{str(run_id)}]: "
                 f"Finished {str(last_run['finished_at_humanized'])} (duration: {str(last_run['duration_humanized'])})"
             )
             try:
@@ -243,6 +259,7 @@ def get_blobs_grouped_by_dir(blobs: List[str]) -> Dict[str, List[str]]:
     return blob_grouped_by_directory
 
 
+# pylint: disable=too-many-locals, too-many-branches
 def download_dbt_files(
     blob_grouped_by_directory: Dict, config, client, bucket_name: Optional[str]
 ) -> Iterable[DbtFiles]:
@@ -255,6 +272,7 @@ def download_dbt_files(
     ) in blob_grouped_by_directory.items():
         dbt_catalog = None
         dbt_manifest = None
+        dbt_sources = None
         dbt_run_results = []
         kwargs = {}
         if bucket_name:
@@ -285,12 +303,16 @@ def download_dbt_files(
                             logger.warning(
                                 f"{DBT_RUN_RESULTS_FILE_NAME} not found in {key}: {exc}"
                             )
+                    if DBT_SOURCES_FILE_NAME == blob_file_name.lower():
+                        logger.debug(f"{DBT_SOURCES_FILE_NAME} found in {key}")
+                        dbt_sources = reader.read(path=blob, **kwargs)
             if not dbt_manifest:
                 raise DBTConfigException(f"Manifest file not found at: {key}")
             yield DbtFiles(
                 dbt_catalog=json.loads(dbt_catalog) if dbt_catalog else None,
                 dbt_manifest=json.loads(dbt_manifest),
                 dbt_run_results=dbt_run_results if dbt_run_results else None,
+                dbt_sources=json.loads(dbt_sources) if dbt_sources else None,
             )
         except DBTConfigException as exc:
             logger.warning(exc)

metadata/ingestion/source/database/dbt/dbt_service.py

@@ -15,7 +15,12 @@ DBT service Topology.
 from abc import ABC, abstractmethod
 from typing import Iterable, List
 
-from dbt_artifacts_parser.parser import parse_catalog, parse_manifest, parse_run_results
+from dbt_artifacts_parser.parser import (
+    parse_catalog,
+    parse_manifest,
+    parse_run_results,
+    parse_sources,
+)
 from pydantic import Field
 from typing_extensions import Annotated
 
@@ -38,6 +43,7 @@ from metadata.ingestion.models.topology import (
 )
 from metadata.ingestion.source.database.database_service import DataModelLink
 from metadata.ingestion.source.database.dbt.constants import (
+    REQUIRED_CONSTRAINT_KEYS,
     REQUIRED_NODE_KEYS,
     REQUIRED_RESULTS_KEYS,
 )
@@ -182,6 +188,20 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
                 ]
                 for key in keys_to_delete:
                     del value[key]
+                if value.get("columns"):
+                    for col_name, value in value[
+                        "columns"
+                    ].items():  # pylint: disable=unused-variable
+                        if value.get("constraints"):
+                            keys_to_delete = [
+                                key
+                                for key in value
+                                if key.lower() not in REQUIRED_CONSTRAINT_KEYS
+                            ]
+                            for key in keys_to_delete:
+                                del value[key]
+                        else:
+                            value["constraints"] = None
 
     def remove_run_result_non_required_keys(self, run_results: List[dict]):
         """
@@ -209,11 +229,15 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
             self.remove_run_result_non_required_keys(
                 run_results=self.context.get().dbt_file.dbt_run_results
             )
+
         dbt_objects = DbtObjects(
             dbt_catalog=parse_catalog(self.context.get().dbt_file.dbt_catalog)
             if self.context.get().dbt_file.dbt_catalog
             else None,
             dbt_manifest=parse_manifest(self.context.get().dbt_file.dbt_manifest),
+            dbt_sources=parse_sources(self.context.get().dbt_file.dbt_sources)
+            if self.context.get().dbt_file.dbt_sources
+            else None,
             dbt_run_results=[
                 parse_run_results(run_result_file)
                 for run_result_file in self.context.get().dbt_file.dbt_run_results

metadata/ingestion/source/database/dbt/dbt_utils.py

@@ -44,6 +44,20 @@ def create_test_case_parameter_definitions(dbt_test):
                 }
             ]
             return test_case_param_definition
+        if hasattr(dbt_test, "freshness"):
+            test_case_param_definition = [
+                {
+                    "name": "warn_after",
+                    "displayName": "warn_after",
+                    "required": False,
+                },
+                {
+                    "name": "error_after",
+                    "displayName": "error_after",
+                    "required": False,
+                },
+            ]
+            return test_case_param_definition
     except Exception as err:  # pylint: disable=broad-except
         logger.debug(traceback.format_exc())
         logger.error(
@@ -67,6 +81,21 @@ def create_test_case_parameter_values(dbt_test):
                 {"name": manifest_node.test_metadata.name, "value": dbt_test_values}
             ]
             return test_case_param_values
+        if hasattr(manifest_node, "freshness"):
+            warn_after = manifest_node.freshness.warn_after
+            error_after = manifest_node.freshness.error_after
+
+            test_case_param_values = [
+                {
+                    "name": "error_after",
+                    "value": f"{error_after.count} {error_after.period.value}",
+                },
+                {
+                    "name": "warn_after",
+                    "value": f"{warn_after.count} {warn_after.period.value}",
+                },
+            ]
+            return test_case_param_values
     except Exception as err:  # pylint: disable=broad-except
         logger.debug(traceback.format_exc())
         logger.error(