openmetadata-ingestion 1.4.6.0rc1__py3-none-any.whl → 1.5.0.0rc1__py3-none-any.whl

_openmetadata_testutils/helpers/assumption.py

@@ -0,0 +1,118 @@
+from typing import Callable, List
+
+import pytest
+from pandas import DataFrame, Series
+
+
+class Assumption:
+    def __init__(self):
+        self.not_ = False
+
+    def validate(self, df: DataFrame):
+        if self.not_:
+            self.assume_false(df)
+        else:
+            self.assume_positive(df)
+
+    def __invert__(self):
+        return self.copy(not_=not self.not_)
+
+    def assume_positive(self, df: DataFrame):
+        raise NotImplementedError
+
+    def assume_false(self, df: DataFrame):
+        with pytest.raises(AssertionError):
+            self.assume_positive(df)
+
+    def copy(self, **kwargs):
+        new = self.__class__.__new__(self.__class__)
+        new.__dict__.update(self.__dict__)
+        new.__dict__.update(kwargs)
+        return new
+
+
+class AssumeUnqiue(Assumption):
+    def __init__(self, column: str):
+        super().__init__()
+        self.column = column
+
+    def assume_positive(self, df: DataFrame):
+        assert df[self.column].nunique() == len(df)
+
+
+class AssumeNotNull(Assumption):
+    def __init__(self, column: str):
+        super().__init__()
+        self.column = column
+
+    def assume_positive(self, df: DataFrame):
+        assert df[self.column].notnull().all()
+
+
+class AssumeMatchRegex(Assumption):
+    def __init__(self, column: str, regex: str):
+        super().__init__()
+        self.column = column
+        self.regex = regex
+
+    def assume_positive(self, df: DataFrame):
+        assert df[self.column].str.match(self.regex).all()
+
+
+class AssumeBetween(Assumption):
+    def __init__(self, column: str, min_value: int, max_value: int):
+        super().__init__()
+        self.column = column
+        self.min_value = min_value
+        self.max_value = max_value
+
+    def assume_positive(self, df: DataFrame):
+        assert (df[self.column] >= self.min_value).all()
+        assert (df[self.column] <= self.max_value).all()
+
+
+class AssumeLengthBetween(Assumption):
+    def __init__(self, column: str, min_length: int, max_length: int):
+        super().__init__()
+        self.column = column
+        self.min_length = min_length
+        self.max_length = max_length
+
+    def assume_positive(self, df: DataFrame):
+        assert (df[self.column].str.len() >= self.min_length).all()
+        assert (df[self.column].str.len() <= self.max_length).all()
+
+
+class AssumeColumnValuesIn(Assumption):
+    def __init__(self, column: str, allowed_values: List[str]):
+        super().__init__()
+        self.column = column
+        self.allowed_values = allowed_values
+
+    def assume_positive(self, df: DataFrame):
+        assert df[self.column].isin(self.allowed_values).all()
+
+
+class AssumeArbitrary(Assumption):
+    def __init__(self, column: str, fn: Callable[[Series], Series]):
+        super().__init__()
+        self.column = column
+        self.fn = fn
+
+    def assume_positive(self, df: DataFrame):
+        assert self.fn(
+            df[self.column]
+        ).all(), f"failed test {self.__class__.__name__} for column {self.column}"
+
+
+class Assumptions:
+    notnull = AssumeNotNull
+    unique = AssumeUnqiue
+    match_regex = AssumeMatchRegex
+    between = AssumeBetween
+    length_between = AssumeLengthBetween
+    column_values_in = AssumeColumnValuesIn
+    arbitrary = AssumeArbitrary
+
+
+assume = Assumptions()

_openmetadata_testutils/helpers/docker.py

@@ -0,0 +1,38 @@
+import contextlib
+import logging
+import os
+import tarfile
+
+import docker
+from docker.models.containers import Container
+
+
+@contextlib.contextmanager
+def try_bind(container, container_port, host_port):
+    """Try to bind a port to the container, if it is already in use try another port."""
+    try:
+        with container.with_bind_ports(container_port, host_port) as container:
+            yield container
+    except docker.errors.APIError:
+        logging.warning("Port %s is already in use, trying another port", host_port)
+        with container.with_bind_ports(container_port, None) as container:
+            yield container
+
+
+def copy_dir_to_container(dir_path: str, container: Container, container_path: str):
+    """Copy the contents of a directory to a path in a container. If the path
+    does not exist it will be created.
+
+    Args:
+        dir_path (str): Path to the directory to copy.
+        container (DockerContainer): The container to copy the directory to.
+        container_path (str): The path to copy the directory to in the container.
+    """
+    tar_path = dir_path + ".tar"
+    with tarfile.open(tar_path, "w") as tar:
+        for item in os.listdir(dir_path):
+            tar.add(os.path.join(dir_path, item), arcname=item)
+    container.exec_run(["mkdir", "-p", container_path])
+    with open(tar_path, "rb") as tar_file:
+        container.put_archive(container_path, tar_file)
+    os.remove(tar_path)

_openmetadata_testutils/helpers/markers.py

@@ -0,0 +1,5 @@
+import pytest
+
+
+def xfail_param(param, reason):
+    return pytest.param(param, marks=pytest.mark.xfail(reason=reason, strict=True))

_openmetadata_testutils/ometa.py

@@ -0,0 +1,23 @@
+from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
+    AuthProvider,
+    OpenMetadataConnection,
+)
+from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
+    OpenMetadataJWTClientConfig,
+)
+from metadata.ingestion.models.custom_pydantic import CustomSecretStr
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+
+OM_JWT = "eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg"
+
+
+def int_admin_ometa(url: str = "http://localhost:8585/api") -> OpenMetadata:
+    """Initialize the ometa connection with default admin:admin creds"""
+    server_config = OpenMetadataConnection(
+        hostPort=url,
+        authProvider=AuthProvider.openmetadata,
+        securityConfig=OpenMetadataJWTClientConfig(jwtToken=CustomSecretStr(OM_JWT)),
+    )
+    metadata = OpenMetadata(server_config)
+    assert metadata.health_check()
+    return metadata

_openmetadata_testutils/postgres/conftest.py

@@ -0,0 +1,70 @@
+"""Test fixtures for PostgreSQL. It uses a testcontainer to start a PostgreSQL instance with the dvdrental database."""
+
+import logging
+import os
+import zipfile
+from subprocess import CalledProcessError
+
+import pytest
+from testcontainers.postgres import PostgresContainer
+
+from _openmetadata_testutils.helpers.docker import copy_dir_to_container, try_bind
+
+
+@pytest.fixture(autouse=True, scope="session")
+def config_logging():
+    """Set the log level of sqlfluff to CRITICAL."""
+    logging.getLogger("sqlfluff").setLevel(logging.CRITICAL)
+
+
+@pytest.fixture(scope="module")
+def postgres_container(tmp_path_factory):
+    """Start a PostgreSQL container with the dvdrental database."""
+    data_dir = tmp_path_factory.mktemp("data")
+    dvd_rental_zip = os.path.join(
+        os.path.dirname(os.path.dirname(__file__)), "data", "dvdrental.zip"
+    )
+    zipfile.ZipFile(dvd_rental_zip, "r").extractall(str(data_dir))
+    container = PostgresContainer("postgres:15", dbname="dvdrental")
+    container._command = [
+        "-c",
+        "shared_preload_libraries=pg_stat_statements",
+        "-c",
+        "pg_stat_statements.max=10000",
+        "-c",
+        "pg_stat_statements.track=all",
+        "-c",
+        "track_commit_timestamp=on",
+    ]
+
+    with (
+        try_bind(container, 5432, 5432) if not os.getenv("CI") else container
+    ) as container:
+        docker_container = container.get_wrapped_container()
+        copy_dir_to_container(str(data_dir), docker_container, "/data")
+        for query in (
+            "CREATE USER postgres SUPERUSER;",
+            "CREATE EXTENSION pg_stat_statements;",
+        ):
+            res = docker_container.exec_run(
+                ["psql", "-U", container.username, "-d", container.dbname, "-c", query]
+            )
+            if res[0] != 0:
+                raise CalledProcessError(
+                    returncode=res[0], cmd=res, output=res[1].decode("utf-8")
+                )
+        res = docker_container.exec_run(
+            [
+                "pg_restore",
+                "-U",
+                container.username,
+                "-d",
+                container.dbname,
+                "/data/dvdrental.tar",
+            ]
+        )
+        if res[0] != 0:
+            raise CalledProcessError(
+                returncode=res[0], cmd=res, output=res[1].decode("utf-8")
+            )
+        yield container

_openmetadata_testutils/pydantic/test_utils.py

@@ -0,0 +1,78 @@
+from collections import deque
+
+from pydantic import BaseModel
+
+
+def assert_equal_pydantic_objects(
+    expected: BaseModel, actual: BaseModel, ignore_none=True
+):
+    """Compare 2 pydantic objects recursively and raise an AssertionError if they are not equal along with all
+    the differences by field. If `ignore_none` is set to True, expected None values will be ignored. This can be
+    useful when comparing objects that are partially filled.
+
+    Example:
+        >>> from pydantic import BaseModel
+        >>> class A(BaseModel):
+        ...     a: int
+        >>> class B(BaseModel):
+        ...     b: A
+        >>> a1 = A(a=1)
+        >>> a2 = A(a=2)
+        >>> b1 = B(b=a1)
+        >>> b2 = B(b=a2)
+        >>> assert_equal_pydantic_objects(a1, b1)
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on type at : expected: [A], actual: [B]
+        >>> assert_equal_pydantic_objects(a1, a2)
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on field: [a], expected: [1], actual: [2]
+        >>> assert_equal_pydantic_objects(b1, b2)
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on field: [b.a], expected: [1], actual: [2]
+
+    Args:
+        expected (BaseModel): The expected pydantic object.
+        actual (BaseModel): The actual pydantic object.
+        ignore_none (bool, optional): Whether to ignore None values. Defaults to True.
+
+    Raises:
+        AssertionError: If the objects are not equal. The error message will contain all the differences.
+    """
+    errors = []
+    queue = deque([(expected, actual, "")])
+
+    while queue:
+        expected, actual, current_key_prefix = queue.popleft()
+        if not isinstance(expected, actual.__class__):
+            errors.append(
+                f"objects mismatched on type at {current_key_prefix}: "
+                f"expected: [{type(expected).__name__}], actual: [{type(actual).__name__}]"
+            )
+            continue
+        if issubclass(expected.__class__, BaseModel):
+            for key, expected_value in expected.dict().items():
+                if expected_value is None and ignore_none:
+                    continue
+                actual_value = actual.dict().get(key)
+                new_key_prefix = (
+                    f"{current_key_prefix}.{key}" if current_key_prefix else key
+                )
+                if issubclass(getattr(expected, key).__class__, BaseModel):
+                    queue.append(
+                        (getattr(expected, key), getattr(actual, key), new_key_prefix)
+                    )
+                elif expected_value != actual_value:
+                    errors.append(
+                        f"objects mismatched on field: [{new_key_prefix}], expected: [{expected_value}], actual: [{actual_value}]"
+                    )
+        else:
+            if expected != actual:
+                errors.append(
+                    f"mismatch at {current_key_prefix}: expected: [{expected}], actual: [{actual}]"
+                )
+
+    if errors:
+        raise AssertionError("\n".join(errors))

airflow_provider_openmetadata/lineage/callback.py

@@ -48,7 +48,7 @@ def failure_callback(context: Dict[str, str]) -> None:
         )
         pipeline: Pipeline = metadata.get_by_name(
             entity=Pipeline,
-            fqn=f"{airflow_service_entity.name.__root__}.{dag.dag_id}",
+            fqn=f"{airflow_service_entity.name.root}.{dag.dag_id}",
         )
 
         if pipeline:
@@ -60,7 +60,7 @@ def failure_callback(context: Dict[str, str]) -> None:
             )
         else:
             logging.warning(
-                f"Pipeline {airflow_service_entity.name.__root__}.{dag.dag_id} not found. Skipping status update."
+                f"Pipeline {airflow_service_entity.name.root}.{dag.dag_id} not found. Skipping status update."
             )
 
     except Exception as exc:  # pylint: disable=broad-except
@@ -90,7 +90,7 @@ def success_callback(context: Dict[str, str]) -> None:
         )
         pipeline: Pipeline = metadata.get_by_name(
             entity=Pipeline,
-            fqn=f"{airflow_service_entity.name.__root__}.{dag.dag_id}",
+            fqn=f"{airflow_service_entity.name.root}.{dag.dag_id}",
         )
 
         add_status(

airflow_provider_openmetadata/lineage/config/loader.py

@@ -90,7 +90,7 @@ def get_lineage_config() -> AirflowLineageConfig:
     if openmetadata_config_file:
         with open(openmetadata_config_file, encoding="utf-8") as config_file:
             config = json.load(config_file)
-            return AirflowLineageConfig.parse_obj(config)
+            return AirflowLineageConfig.model_validate(config)
 
     # If nothing is configured, raise
     raise ValueError("Missing lineage backend configuration")

airflow_provider_openmetadata/lineage/runner.py

@@ -200,7 +200,7 @@ class AirflowLineageRunner:
             return self.metadata.patch(
                 entity=Pipeline,
                 source=pipeline,
-                destination=pipeline.copy(update=pipeline_request.__dict__),
+                destination=pipeline.model_copy(update=pipeline_request.__dict__),
                 allowed_fields=ALLOWED_COMMON_PATCH_FIELDS,
                 restrict_update_fields=RESTRICT_UPDATE_LIST,
             )
@@ -282,7 +282,7 @@ class AirflowLineageRunner:
 
         for status in pipeline_status_list:
             self.metadata.add_pipeline_status(
-                fqn=pipeline.fullyQualifiedName.__root__, status=status
+                fqn=pipeline.fullyQualifiedName.root, status=status
             )
 
     def add_lineage(self, pipeline: Pipeline, xlets: XLets) -> None:
@@ -327,12 +327,12 @@ class AirflowLineageRunner:
                     else:
                         self.dag.log.warning(
                             f"Could not find [{to_xlet.entity.__name__}] [{to_xlet.fqn}] from "
-                            f"[{pipeline.fullyQualifiedName.__root__}] outlets"
+                            f"[{pipeline.fullyQualifiedName.root}] outlets"
                         )
             else:
                 self.dag.log.warning(
                     f"Could not find [{from_xlet.entity.__name__}] [{from_xlet.fqn}] from "
-                    f"[{pipeline.fullyQualifiedName.__root__}] inlets"
+                    f"[{pipeline.fullyQualifiedName.root}] inlets"
                 )
 
     def clean_lineage(self, pipeline: Pipeline, xlets: XLets):
@@ -343,7 +343,7 @@ class AirflowLineageRunner:
         """
         lineage_data = self.metadata.get_lineage_by_name(
             entity=Pipeline,
-            fqn=pipeline.fullyQualifiedName.__root__,
+            fqn=pipeline.fullyQualifiedName.root,
             up_depth=1,
             down_depth=1,
         )

airflow_provider_openmetadata/lineage/status.py

@@ -91,7 +91,7 @@ def add_status(
 
     task_status = []
     # We will append based on the current registered status
-    if pipeline_status and pipeline_status.timestamp.__root__ == execution_date:
+    if pipeline_status and pipeline_status.timestamp.root == execution_date:
         # If we are clearing a task, use the status of the new execution
         task_status = [
             task
@@ -123,5 +123,5 @@ def add_status(
 
     operator.log.info(f"Added status to DAG {updated_status}")
     metadata.add_pipeline_status(
-        fqn=pipeline.fullyQualifiedName.__root__, status=updated_status
+        fqn=pipeline.fullyQualifiedName.root, status=updated_status
     )

metadata/automations/runner.py

@@ -37,7 +37,7 @@ def execute(encrypted_automation_workflow: AutomationWorkflow) -> Any:
     )
 
     automation_workflow = metadata.get_by_name(
-        entity=AutomationWorkflow, fqn=encrypted_automation_workflow.name.__root__
+        entity=AutomationWorkflow, fqn=encrypted_automation_workflow.name.root
     )
 
     return run_workflow(automation_workflow.request, automation_workflow, metadata)

metadata/cli/app.py

@@ -19,7 +19,6 @@ from pathlib import Path
 from metadata.config.common import load_config_file
 from metadata.utils.logger import cli_logger
 from metadata.workflow.application import ApplicationWorkflow
-from metadata.workflow.application_output_handler import print_status
 
 logger = cli_logger()
 
@@ -42,5 +41,5 @@ def run_app(config_path: Path) -> None:
 
     workflow.execute()
     workflow.stop()
-    print_status(workflow)
+    workflow.print_status()
     workflow.raise_from_status()

metadata/cli/dataquality.py

@@ -17,13 +17,12 @@ import traceback
 from pathlib import Path
 
 from metadata.config.common import load_config_file
+from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
+    PipelineType,
+)
 from metadata.utils.logger import cli_logger
 from metadata.workflow.data_quality import TestSuiteWorkflow
-from metadata.workflow.workflow_output_handler import (
-    WorkflowType,
-    print_init_error,
-    print_status,
-)
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
 
@@ -42,10 +41,12 @@ def run_test(config_path: Path) -> None:
         workflow = TestSuiteWorkflow.create(workflow_config_dict)
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        print_init_error(exc, workflow_config_dict, WorkflowType.TEST)
+        WorkflowInitErrorHandler.print_init_error(
+            exc, workflow_config_dict, PipelineType.TestSuite
+        )
         sys.exit(1)
 
     workflow.execute()
     workflow.stop()
-    print_status(workflow)
+    workflow.print_status()
     workflow.raise_from_status()

metadata/cli/ingest.py

@@ -17,13 +17,12 @@ import traceback
 from pathlib import Path
 
 from metadata.config.common import load_config_file
+from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
+    PipelineType,
+)
 from metadata.utils.logger import cli_logger
 from metadata.workflow.metadata import MetadataWorkflow
-from metadata.workflow.workflow_output_handler import (
-    WorkflowType,
-    print_init_error,
-    print_status,
-)
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
 
@@ -42,10 +41,12 @@ def run_ingest(config_path: Path) -> None:
         logger.debug(f"Using config: {workflow.config}")
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        print_init_error(exc, config_dict, WorkflowType.INGEST)
+        WorkflowInitErrorHandler.print_init_error(
+            exc, config_dict, PipelineType.metadata
+        )
         sys.exit(1)
 
     workflow.execute()
     workflow.stop()
-    print_status(workflow)
+    workflow.print_status()
     workflow.raise_from_status()

metadata/cli/lineage.py

@@ -21,18 +21,21 @@ from pydantic import BaseModel
 
 from metadata.config.common import load_config_file
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
+    PipelineType,
+)
 from metadata.generated.schema.metadataIngestion.workflow import WorkflowConfig
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.constants import UTF_8
 from metadata.utils.logger import cli_logger
-from metadata.workflow.workflow_output_handler import WorkflowType, print_init_error
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
 
 
 class LineageWorkflow(BaseModel):
-    filePath: Optional[str]
-    query: Optional[str]
+    filePath: Optional[str] = None
+    query: Optional[str] = None
     checkPatch: Optional[bool] = True
     serviceName: str
     workflowConfig: WorkflowConfig
@@ -49,11 +52,13 @@ def run_lineage(config_path: Path) -> None:
     config_dict = None
     try:
         config_dict = load_config_file(config_path)
-        workflow = LineageWorkflow.parse_obj(config_dict)
+        workflow = LineageWorkflow.model_validate(config_dict)
 
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        print_init_error(exc, config_dict, WorkflowType.INGEST)
+        WorkflowInitErrorHandler.print_init_error(
+            exc, config_dict, PipelineType.lineage
+        )
         sys.exit(1)
 
     if workflow.filePath:

metadata/cli/profile.py

@@ -17,13 +17,12 @@ import traceback
 from pathlib import Path
 
 from metadata.config.common import load_config_file
+from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
+    PipelineType,
+)
 from metadata.utils.logger import cli_logger
 from metadata.workflow.profiler import ProfilerWorkflow
-from metadata.workflow.workflow_output_handler import (
-    WorkflowType,
-    print_init_error,
-    print_status,
-)
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
 
@@ -42,10 +41,12 @@ def run_profiler(config_path: Path) -> None:
         workflow = ProfilerWorkflow.create(workflow_config_dict)
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        print_init_error(exc, workflow_config_dict, WorkflowType.PROFILE)
+        WorkflowInitErrorHandler.print_init_error(
+            exc, workflow_config_dict, PipelineType.profiler
+        )
         sys.exit(1)
 
     workflow.execute()
     workflow.stop()
-    print_status(workflow)
+    workflow.print_status()
     workflow.raise_from_status()

metadata/cli/usage.py

@@ -17,13 +17,12 @@ import traceback
 from pathlib import Path
 
 from metadata.config.common import load_config_file
+from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
+    PipelineType,
+)
 from metadata.utils.logger import cli_logger
 from metadata.workflow.usage import UsageWorkflow
-from metadata.workflow.workflow_output_handler import (
-    WorkflowType,
-    print_init_error,
-    print_status,
-)
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
 
@@ -42,10 +41,10 @@ def run_usage(config_path: Path) -> None:
         logger.debug(f"Using config: {workflow.config}")
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        print_init_error(exc, config_dict, WorkflowType.INGEST)
+        WorkflowInitErrorHandler.print_init_error(exc, config_dict, PipelineType.usage)
         sys.exit(1)
 
     workflow.execute()
     workflow.stop()
-    print_status(workflow)
+    workflow.print_status()
     workflow.raise_from_status()

metadata/clients/aws_client.py

@@ -47,7 +47,7 @@ class AWSAssumeRoleException(Exception):
 class AWSAssumeRoleCredentialWrapper(BaseModel):
     accessKeyId: str
     secretAccessKey: CustomSecretStr
-    sessionToken: Optional[str]
+    sessionToken: Optional[str] = None
 
 
 class AWSClient:
@@ -59,7 +59,7 @@ class AWSClient:
         self.config = (
             config
             if isinstance(config, AWSCredentials)
-            else (AWSCredentials.parse_obj(config) if config else config)
+            else (AWSCredentials.model_validate(config) if config else config)
         )
 
     @staticmethod
@@ -148,7 +148,7 @@ class AWSClient:
             session = self.create_session()
             if self.config.endPointURL is not None:
                 return session.client(
-                    service_name=service_name, endpoint_url=self.config.endPointURL
+                    service_name=service_name, endpoint_url=str(self.config.endPointURL)
                 )
             return session.client(service_name=service_name)
 
@@ -160,7 +160,7 @@ class AWSClient:
         session = self.create_session()
         if self.config.endPointURL is not None:
             return session.resource(
-                service_name=service_name, endpoint_url=self.config.endPointURL
+                service_name=service_name, endpoint_url=str(self.config.endPointURL)
             )
         return session.resource(service_name=service_name)