From 58f483de73a9374511a69720f46475426b2af398 Mon Sep 17 00:00:00 2001 From: Justin Bandoro Date: Tue, 27 Feb 2024 10:44:33 -0800 Subject: [PATCH] add ruff import sorting (isort) --- cosmos/__init__.py | 6 +-- cosmos/airflow/dag.py | 2 +- cosmos/airflow/graph.py | 12 ++--- cosmos/airflow/task_group.py | 3 +- cosmos/config.py | 8 +-- cosmos/constants.py | 1 - cosmos/converter.py | 4 +- cosmos/core/airflow.py | 1 - cosmos/dbt/graph.py | 5 +- cosmos/dbt/parser/output.py | 3 +- cosmos/dbt/parser/project.py | 2 +- cosmos/dbt/project.py | 7 +-- cosmos/dbt/selector.py | 1 + cosmos/hooks/subprocess.py | 2 +- cosmos/log.py | 2 +- cosmos/operators/__init__.py | 2 +- cosmos/operators/azure_container_instance.py | 6 +-- cosmos/operators/base.py | 3 +- cosmos/operators/docker.py | 4 +- cosmos/operators/kubernetes.py | 11 ++--- cosmos/operators/local.py | 49 +++++++++---------- cosmos/operators/virtualenv.py | 4 +- cosmos/profiles/__init__.py | 7 ++- cosmos/profiles/base.py | 6 +-- cosmos/profiles/bigquery/__init__.py | 2 +- .../bigquery/service_account_keyfile_dict.py | 4 +- cosmos/profiles/snowflake/__init__.py | 4 +- dev/dags/basic_cosmos_dag.py | 2 +- dev/dags/basic_cosmos_task_group.py | 5 +- dev/dags/cosmos_manifest_example.py | 4 +- dev/dags/cosmos_profile_mapping.py | 2 +- dev/dags/dbt/simple/convert_csv_to_db.py | 3 +- dev/dags/dbt_docs.py | 6 +-- dev/dags/example_cosmos_python_models.py | 2 +- dev/dags/example_cosmos_sources.py | 4 +- dev/dags/example_model_version.py | 2 +- dev/dags/example_virtualenv.py | 2 +- dev/dags/performance_dag.py | 5 +- dev/dags/user_defined_profile.py | 2 +- docs/generate_mappings.py | 3 +- pyproject.toml | 2 +- tests/airflow/test_graph.py | 2 +- tests/dbt/parser/test_output.py | 5 +- tests/dbt/test_graph.py | 8 +-- tests/dbt/test_project.py | 2 +- tests/dbt/test_selector.py | 3 +- tests/hooks/test_subprocess.py | 2 +- .../test_azure_container_instance.py | 2 +- tests/operators/test_base.py | 9 ++-- tests/operators/test_docker.py | 2 +- tests/operators/test_kubernetes.py | 5 +- tests/operators/test_local.py | 32 ++++++------ tests/operators/test_virtualenv.py | 8 ++- tests/perf/test_performance.py | 4 +- tests/plugin/test_plugin.py | 12 ++--- .../profiles/athena/test_athena_access_key.py | 3 +- .../test_bq_service_account_keyfile_dict.py | 2 +- tests/profiles/test_base_profile.py | 3 +- tests/profiles/trino/test_trino_base.py | 2 +- tests/profiles/trino/test_trino_jwt.py | 3 +- tests/profiles/trino/test_trino_ldap.py | 3 +- tests/test_config.py | 7 ++- tests/test_converter.py | 9 ++-- tests/test_example_dags_no_connections.py | 1 - tests/test_log.py | 3 +- 65 files changed, 163 insertions(+), 174 deletions(-) diff --git a/cosmos/__init__.py b/cosmos/__init__.py index c5a0ce845..859995fe6 100644 --- a/cosmos/__init__.py +++ b/cosmos/__init__.py @@ -11,12 +11,12 @@ from cosmos.airflow.dag import DbtDag from cosmos.airflow.task_group import DbtTaskGroup from cosmos.config import ( - ProjectConfig, - ProfileConfig, ExecutionConfig, + ProfileConfig, + ProjectConfig, RenderConfig, ) -from cosmos.constants import LoadMode, TestBehavior, ExecutionMode +from cosmos.constants import ExecutionMode, LoadMode, TestBehavior from cosmos.log import get_logger from cosmos.operators.lazy_load import MissingPackage from cosmos.operators.local import ( diff --git a/cosmos/airflow/dag.py b/cosmos/airflow/dag.py index ca2672060..de958f118 100644 --- a/cosmos/airflow/dag.py +++ b/cosmos/airflow/dag.py @@ -8,7 +8,7 @@ from airflow.models.dag import DAG -from cosmos.converter import airflow_kwargs, specific_kwargs, DbtToAirflowConverter +from cosmos.converter import DbtToAirflowConverter, airflow_kwargs, specific_kwargs class DbtDag(DAG, DbtToAirflowConverter): diff --git a/cosmos/airflow/graph.py b/cosmos/airflow/graph.py index 177806480..905c845a6 100644 --- a/cosmos/airflow/graph.py +++ b/cosmos/airflow/graph.py @@ -1,26 +1,24 @@ from __future__ import annotations -from typing import Any, Callable +from typing import Any, Callable, Union from airflow.models import BaseOperator from airflow.models.dag import DAG from airflow.utils.task_group import TaskGroup +from cosmos.config import RenderConfig from cosmos.constants import ( + DEFAULT_DBT_RESOURCES, + TESTABLE_DBT_RESOURCES, DbtResourceType, + ExecutionMode, TestBehavior, TestIndirectSelection, - ExecutionMode, - TESTABLE_DBT_RESOURCES, - DEFAULT_DBT_RESOURCES, ) -from cosmos.config import RenderConfig from cosmos.core.airflow import get_airflow_task as create_airflow_task from cosmos.core.graph.entities import Task as TaskMetadata from cosmos.dbt.graph import DbtNode from cosmos.log import get_logger -from typing import Union - logger = get_logger(__name__) diff --git a/cosmos/airflow/task_group.py b/cosmos/airflow/task_group.py index 5171645f2..64fcb298a 100644 --- a/cosmos/airflow/task_group.py +++ b/cosmos/airflow/task_group.py @@ -3,11 +3,12 @@ """ from __future__ import annotations + from typing import Any from airflow.utils.task_group import TaskGroup -from cosmos.converter import airflow_kwargs, specific_kwargs, DbtToAirflowConverter +from cosmos.converter import DbtToAirflowConverter, airflow_kwargs, specific_kwargs class DbtTaskGroup(TaskGroup, DbtToAirflowConverter): diff --git a/cosmos/config.py b/cosmos/config.py index 9dfb672be..729e95c75 100644 --- a/cosmos/config.py +++ b/cosmos/config.py @@ -5,18 +5,18 @@ import contextlib import shutil import tempfile +import warnings from dataclasses import InitVar, dataclass, field from pathlib import Path -import warnings -from typing import Any, Iterator, Callable +from typing import Any, Callable, Iterator from cosmos.constants import ( DbtResourceType, - TestBehavior, ExecutionMode, + InvocationMode, LoadMode, + TestBehavior, TestIndirectSelection, - InvocationMode, ) from cosmos.dbt.executable import get_system_dbt from cosmos.exceptions import CosmosValueError diff --git a/cosmos/constants.py b/cosmos/constants.py index e8b9cff1d..1db78d15b 100644 --- a/cosmos/constants.py +++ b/cosmos/constants.py @@ -4,7 +4,6 @@ import aenum - DBT_PROFILE_PATH = Path(os.path.expanduser("~")).joinpath(".dbt/profiles.yml") DEFAULT_DBT_PROFILE_NAME = "cosmos_profile" DEFAULT_DBT_TARGET_NAME = "cosmos_target" diff --git a/cosmos/converter.py b/cosmos/converter.py index bafe094e8..3619fecf5 100644 --- a/cosmos/converter.py +++ b/cosmos/converter.py @@ -3,19 +3,19 @@ from __future__ import annotations +import copy import inspect from typing import Any, Callable -import copy from warnings import warn from airflow.models.dag import DAG from airflow.utils.task_group import TaskGroup from cosmos.airflow.graph import build_airflow_graph +from cosmos.config import ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig from cosmos.constants import ExecutionMode from cosmos.dbt.graph import DbtGraph from cosmos.dbt.selector import retrieve_by_label -from cosmos.config import ProjectConfig, ExecutionConfig, RenderConfig, ProfileConfig from cosmos.exceptions import CosmosValueError from cosmos.log import get_logger diff --git a/cosmos/core/airflow.py b/cosmos/core/airflow.py index 7c5dee328..f6f7464d8 100644 --- a/cosmos/core/airflow.py +++ b/cosmos/core/airflow.py @@ -7,7 +7,6 @@ from cosmos.core.graph.entities import Task from cosmos.log import get_logger - logger = get_logger(__name__) diff --git a/cosmos/dbt/graph.py b/cosmos/dbt/graph.py index 976a736dc..979946f5b 100644 --- a/cosmos/dbt/graph.py +++ b/cosmos/dbt/graph.py @@ -4,12 +4,13 @@ import json import os import tempfile -import yaml from dataclasses import dataclass, field from pathlib import Path from subprocess import PIPE, Popen from typing import Any +import yaml + from cosmos.config import ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig from cosmos.constants import ( DBT_LOG_DIR_NAME, @@ -22,7 +23,7 @@ LoadMode, ) from cosmos.dbt.parser.project import LegacyDbtProject -from cosmos.dbt.project import create_symlinks, copy_msgpack_for_partial_parse, environ +from cosmos.dbt.project import copy_msgpack_for_partial_parse, create_symlinks, environ from cosmos.dbt.selector import select_nodes from cosmos.log import get_logger diff --git a/cosmos/dbt/parser/output.py b/cosmos/dbt/parser/output.py index 3690a8f60..3ff377941 100644 --- a/cosmos/dbt/parser/output.py +++ b/cosmos/dbt/parser/output.py @@ -2,14 +2,13 @@ import logging import re -from typing import List, Tuple, TYPE_CHECKING +from typing import TYPE_CHECKING, List, Tuple if TYPE_CHECKING: from dbt.cli.main import dbtRunnerResult from cosmos.hooks.subprocess import FullOutputSubprocessResult - DBT_NO_TESTS_MSG = "Nothing to do" DBT_WARN_MSG = "WARN" diff --git a/cosmos/dbt/parser/project.py b/cosmos/dbt/parser/project.py index c5e434996..8876d6878 100644 --- a/cosmos/dbt/parser/project.py +++ b/cosmos/dbt/parser/project.py @@ -4,8 +4,8 @@ from __future__ import annotations -import os import ast +import os from dataclasses import dataclass, field from enum import Enum from pathlib import Path diff --git a/cosmos/dbt/project.py b/cosmos/dbt/project.py index 144a1f6df..ad328a332 100644 --- a/cosmos/dbt/project.py +++ b/cosmos/dbt/project.py @@ -1,12 +1,13 @@ from __future__ import annotations -from pathlib import Path -import shutil import os -from cosmos.constants import DBT_LOG_DIR_NAME, DBT_TARGET_DIR_NAME, DBT_PARTIAL_PARSE_FILE_NAME +import shutil from contextlib import contextmanager +from pathlib import Path from typing import Generator +from cosmos.constants import DBT_LOG_DIR_NAME, DBT_PARTIAL_PARSE_FILE_NAME, DBT_TARGET_DIR_NAME + def create_symlinks(project_path: Path, tmp_dir: Path, ignore_dbt_packages: bool) -> None: """Helper function to create symlinks to the dbt project files.""" diff --git a/cosmos/dbt/selector.py b/cosmos/dbt/selector.py index 58c5c12b4..61458c4aa 100644 --- a/cosmos/dbt/selector.py +++ b/cosmos/dbt/selector.py @@ -1,4 +1,5 @@ from __future__ import annotations + import copy import re from collections import defaultdict diff --git a/cosmos/hooks/subprocess.py b/cosmos/hooks/subprocess.py index 2522420b7..19c88540b 100644 --- a/cosmos/hooks/subprocess.py +++ b/cosmos/hooks/subprocess.py @@ -7,9 +7,9 @@ import contextlib import os import signal -from typing import NamedTuple from subprocess import PIPE, STDOUT, Popen from tempfile import TemporaryDirectory, gettempdir +from typing import NamedTuple from airflow.hooks.base import BaseHook diff --git a/cosmos/log.py b/cosmos/log.py index e4ad6e2ba..f7c512f17 100644 --- a/cosmos/log.py +++ b/cosmos/log.py @@ -1,10 +1,10 @@ from __future__ import annotations + import logging from airflow.configuration import conf from airflow.utils.log.colored_log import CustomTTYColoredFormatter - LOG_FORMAT: str = ( "[%(blue)s%(asctime)s%(reset)s] " "{%(blue)s%(filename)s:%(reset)s%(lineno)d} " diff --git a/cosmos/operators/__init__.py b/cosmos/operators/__init__.py index 92f53fa08..c546da019 100644 --- a/cosmos/operators/__init__.py +++ b/cosmos/operators/__init__.py @@ -1,9 +1,9 @@ from .local import DbtBuildLocalOperator as DbtBuildOperator from .local import DbtDepsLocalOperator as DbtDepsOperator from .local import DbtDocsAzureStorageLocalOperator as DbtDocsAzureStorageOperator +from .local import DbtDocsGCSLocalOperator as DbtDocsGCSOperator from .local import DbtDocsLocalOperator as DbtDocsOperator from .local import DbtDocsS3LocalOperator as DbtDocsS3Operator -from .local import DbtDocsGCSLocalOperator as DbtDocsGCSOperator from .local import DbtLSLocalOperator as DbtLSOperator from .local import DbtRunLocalOperator as DbtRunOperator from .local import DbtRunOperationLocalOperator as DbtRunOperationOperator diff --git a/cosmos/operators/azure_container_instance.py b/cosmos/operators/azure_container_instance.py index 903524533..397a47551 100644 --- a/cosmos/operators/azure_container_instance.py +++ b/cosmos/operators/azure_container_instance.py @@ -3,17 +3,17 @@ from typing import Any, Callable, Sequence from airflow.utils.context import Context -from cosmos.config import ProfileConfig +from cosmos.config import ProfileConfig from cosmos.log import get_logger from cosmos.operators.base import ( AbstractDbtBaseOperator, + DbtLSMixin, DbtRunMixin, + DbtRunOperationMixin, DbtSeedMixin, DbtSnapshotMixin, DbtTestMixin, - DbtLSMixin, - DbtRunOperationMixin, ) logger = get_logger(__name__) diff --git a/cosmos/operators/base.py b/cosmos/operators/base.py index d8fefc523..e94cae05f 100644 --- a/cosmos/operators/base.py +++ b/cosmos/operators/base.py @@ -1,8 +1,8 @@ from __future__ import annotations import os -from typing import Any, Sequence, Tuple from abc import ABCMeta, abstractmethod +from typing import Any, Sequence, Tuple import yaml from airflow.models.baseoperator import BaseOperator @@ -12,7 +12,6 @@ from cosmos.dbt.executable import get_system_dbt from cosmos.log import get_logger - logger = get_logger(__name__) diff --git a/cosmos/operators/docker.py b/cosmos/operators/docker.py index 571bff046..5be03fad7 100644 --- a/cosmos/operators/docker.py +++ b/cosmos/operators/docker.py @@ -8,12 +8,12 @@ from cosmos.operators.base import ( AbstractDbtBaseOperator, DbtBuildMixin, + DbtLSMixin, DbtRunMixin, + DbtRunOperationMixin, DbtSeedMixin, DbtSnapshotMixin, DbtTestMixin, - DbtLSMixin, - DbtRunOperationMixin, ) logger = get_logger(__name__) diff --git a/cosmos/operators/kubernetes.py b/cosmos/operators/kubernetes.py index fc0e28c05..e314b7c43 100644 --- a/cosmos/operators/kubernetes.py +++ b/cosmos/operators/kubernetes.py @@ -3,24 +3,23 @@ from os import PathLike from typing import Any, Callable, Sequence +from airflow.models import TaskInstance from airflow.utils.context import Context, context_merge -from cosmos.log import get_logger from cosmos.config import ProfileConfig +from cosmos.dbt.parser.output import extract_log_issues +from cosmos.log import get_logger from cosmos.operators.base import ( AbstractDbtBaseOperator, DbtBuildMixin, + DbtLSMixin, DbtRunMixin, + DbtRunOperationMixin, DbtSeedMixin, DbtSnapshotMixin, DbtTestMixin, - DbtLSMixin, - DbtRunOperationMixin, ) -from airflow.models import TaskInstance -from cosmos.dbt.parser.output import extract_log_issues - DBT_NO_TESTS_MSG = "Nothing to do" DBT_WARN_MSG = "WARN" diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py index fa4db144d..7e71fcba6 100644 --- a/cosmos/operators/local.py +++ b/cosmos/operators/local.py @@ -1,28 +1,28 @@ from __future__ import annotations import os -import signal import tempfile -from attr import define -from pathlib import Path -from typing import Any, Callable, Literal, Sequence, TYPE_CHECKING -from abc import ABC, abstractmethod import warnings +from abc import ABC, abstractmethod +from functools import cached_property +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Literal, Sequence import airflow import jinja2 from airflow import DAG -from functools import cached_property from airflow.configuration import conf from airflow.exceptions import AirflowException, AirflowSkipException from airflow.models.taskinstance import TaskInstance from airflow.utils.context import Context from airflow.utils.session import NEW_SESSION, create_session, provide_session +from attr import define + from cosmos.constants import InvocationMode try: - from openlineage.common.provider.dbt.local import DbtLocalArtifactProcessor from airflow.datasets import Dataset + from openlineage.common.provider.dbt.local import DbtLocalArtifactProcessor except ModuleNotFoundError: is_openlineage_available = False DbtLocalArtifactProcessor = None @@ -31,41 +31,40 @@ if TYPE_CHECKING: from airflow.datasets import Dataset # noqa: F811 - from openlineage.client.run import RunEvent from dbt.cli.main import dbtRunner, dbtRunnerResult + from openlineage.client.run import RunEvent from sqlalchemy.orm import Session +from cosmos.config import ProfileConfig from cosmos.constants import ( + DBT_PARTIAL_PARSE_FILE_NAME, + DBT_TARGET_DIR_NAME, DEFAULT_OPENLINEAGE_NAMESPACE, OPENLINEAGE_PRODUCER, - DBT_TARGET_DIR_NAME, - DBT_PARTIAL_PARSE_FILE_NAME, ) -from cosmos.config import ProfileConfig +from cosmos.dbt.parser.output import ( + extract_dbt_runner_issues, + extract_log_issues, + parse_number_of_warnings_dbt_runner, + parse_number_of_warnings_subprocess, +) +from cosmos.dbt.project import change_working_directory, copy_msgpack_for_partial_parse, create_symlinks, environ +from cosmos.hooks.subprocess import ( + FullOutputSubprocessHook, + FullOutputSubprocessResult, +) from cosmos.log import get_logger from cosmos.operators.base import ( AbstractDbtBaseOperator, DbtBuildMixin, + DbtLSMixin, DbtRunMixin, + DbtRunOperationMixin, DbtSeedMixin, DbtSnapshotMixin, DbtTestMixin, - DbtLSMixin, - DbtRunOperationMixin, ) -from cosmos.hooks.subprocess import ( - FullOutputSubprocessHook, - FullOutputSubprocessResult, -) -from cosmos.dbt.parser.output import ( - extract_dbt_runner_issues, - extract_log_issues, - parse_number_of_warnings_dbt_runner, - parse_number_of_warnings_subprocess, -) -from cosmos.dbt.project import create_symlinks, copy_msgpack_for_partial_parse, environ, change_working_directory - logger = get_logger(__name__) diff --git a/cosmos/operators/virtualenv.py b/cosmos/operators/virtualenv.py index 6612ab8b8..2261f764c 100644 --- a/cosmos/operators/virtualenv.py +++ b/cosmos/operators/virtualenv.py @@ -1,13 +1,13 @@ from __future__ import annotations +from functools import cached_property from pathlib import Path from tempfile import TemporaryDirectory from typing import TYPE_CHECKING, Any -from functools import cached_property from airflow.utils.python_virtualenv import prepare_virtualenv -from cosmos.hooks.subprocess import FullOutputSubprocessResult +from cosmos.hooks.subprocess import FullOutputSubprocessResult from cosmos.log import get_logger from cosmos.operators.local import ( DbtBuildLocalOperator, diff --git a/cosmos/profiles/__init__.py b/cosmos/profiles/__init__.py index 446207f35..fa8e5c370 100644 --- a/cosmos/profiles/__init__.py +++ b/cosmos/profiles/__init__.py @@ -4,20 +4,19 @@ from typing import Any, Type - from .athena import AthenaAccessKeyProfileMapping from .base import BaseProfileMapping, DbtProfileConfigVars +from .bigquery.oauth import GoogleCloudOauthProfileMapping from .bigquery.service_account_file import GoogleCloudServiceAccountFileProfileMapping from .bigquery.service_account_keyfile_dict import GoogleCloudServiceAccountDictProfileMapping -from .bigquery.oauth import GoogleCloudOauthProfileMapping from .databricks.token import DatabricksTokenProfileMapping from .exasol.user_pass import ExasolUserPasswordProfileMapping from .postgres.user_pass import PostgresUserPasswordProfileMapping from .redshift.user_pass import RedshiftUserPasswordProfileMapping +from .snowflake.user_encrypted_privatekey_env_variable import SnowflakeEncryptedPrivateKeyPemProfileMapping +from .snowflake.user_encrypted_privatekey_file import SnowflakeEncryptedPrivateKeyFilePemProfileMapping from .snowflake.user_pass import SnowflakeUserPasswordProfileMapping from .snowflake.user_privatekey import SnowflakePrivateKeyPemProfileMapping -from .snowflake.user_encrypted_privatekey_file import SnowflakeEncryptedPrivateKeyFilePemProfileMapping -from .snowflake.user_encrypted_privatekey_env_variable import SnowflakeEncryptedPrivateKeyPemProfileMapping from .spark.thrift import SparkThriftProfileMapping from .trino.certificate import TrinoCertificateProfileMapping from .trino.jwt import TrinoJWTProfileMapping diff --git a/cosmos/profiles/base.py b/cosmos/profiles/base.py index 1131dc8e1..9d17c4019 100644 --- a/cosmos/profiles/base.py +++ b/cosmos/profiles/base.py @@ -5,13 +5,13 @@ from __future__ import annotations -from abc import ABC, abstractmethod -from typing import Any, Optional, Literal, Dict, TYPE_CHECKING import warnings +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional +import yaml from airflow.hooks.base import BaseHook from pydantic import dataclasses -import yaml from cosmos.exceptions import CosmosValueError from cosmos.log import get_logger diff --git a/cosmos/profiles/bigquery/__init__.py b/cosmos/profiles/bigquery/__init__.py index c03f9a0ab..b547bbf84 100644 --- a/cosmos/profiles/bigquery/__init__.py +++ b/cosmos/profiles/bigquery/__init__.py @@ -1,8 +1,8 @@ "BigQuery Airflow connection -> dbt profile mappings" +from .oauth import GoogleCloudOauthProfileMapping from .service_account_file import GoogleCloudServiceAccountFileProfileMapping from .service_account_keyfile_dict import GoogleCloudServiceAccountDictProfileMapping -from .oauth import GoogleCloudOauthProfileMapping __all__ = [ "GoogleCloudServiceAccountFileProfileMapping", diff --git a/cosmos/profiles/bigquery/service_account_keyfile_dict.py b/cosmos/profiles/bigquery/service_account_keyfile_dict.py index e84587faf..038b34153 100644 --- a/cosmos/profiles/bigquery/service_account_keyfile_dict.py +++ b/cosmos/profiles/bigquery/service_account_keyfile_dict.py @@ -1,10 +1,10 @@ "Maps Airflow GCP connections to dbt BigQuery profiles if they use a service account keyfile dict/json." from __future__ import annotations -from typing import Any import json -from cosmos.exceptions import CosmosValueError +from typing import Any +from cosmos.exceptions import CosmosValueError from cosmos.profiles.base import BaseProfileMapping diff --git a/cosmos/profiles/snowflake/__init__.py b/cosmos/profiles/snowflake/__init__.py index fdf323a76..01745d01e 100644 --- a/cosmos/profiles/snowflake/__init__.py +++ b/cosmos/profiles/snowflake/__init__.py @@ -1,9 +1,9 @@ "Snowflake Airflow connection -> dbt profile mapping." +from .user_encrypted_privatekey_env_variable import SnowflakeEncryptedPrivateKeyPemProfileMapping +from .user_encrypted_privatekey_file import SnowflakeEncryptedPrivateKeyFilePemProfileMapping from .user_pass import SnowflakeUserPasswordProfileMapping from .user_privatekey import SnowflakePrivateKeyPemProfileMapping -from .user_encrypted_privatekey_file import SnowflakeEncryptedPrivateKeyFilePemProfileMapping -from .user_encrypted_privatekey_env_variable import SnowflakeEncryptedPrivateKeyPemProfileMapping __all__ = [ "SnowflakeUserPasswordProfileMapping", diff --git a/dev/dags/basic_cosmos_dag.py b/dev/dags/basic_cosmos_dag.py index 80e3c1a5f..c961638b3 100644 --- a/dev/dags/basic_cosmos_dag.py +++ b/dev/dags/basic_cosmos_dag.py @@ -6,7 +6,7 @@ from datetime import datetime from pathlib import Path -from cosmos import DbtDag, ProjectConfig, ProfileConfig +from cosmos import DbtDag, ProfileConfig, ProjectConfig from cosmos.profiles import PostgresUserPasswordProfileMapping DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" diff --git a/dev/dags/basic_cosmos_task_group.py b/dev/dags/basic_cosmos_task_group.py index 06b24f291..55842ee10 100644 --- a/dev/dags/basic_cosmos_task_group.py +++ b/dev/dags/basic_cosmos_task_group.py @@ -3,16 +3,15 @@ """ import os - from datetime import datetime from pathlib import Path from airflow.decorators import dag from airflow.operators.empty import EmptyOperator -from cosmos import DbtTaskGroup, ProjectConfig, ProfileConfig, RenderConfig, ExecutionConfig -from cosmos.profiles import PostgresUserPasswordProfileMapping +from cosmos import DbtTaskGroup, ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig from cosmos.constants import InvocationMode +from cosmos.profiles import PostgresUserPasswordProfileMapping DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" DBT_ROOT_PATH = Path(os.getenv("DBT_ROOT_PATH", DEFAULT_DBT_ROOT_PATH)) diff --git a/dev/dags/cosmos_manifest_example.py b/dev/dags/cosmos_manifest_example.py index 7b7f9d4aa..e8721aefd 100644 --- a/dev/dags/cosmos_manifest_example.py +++ b/dev/dags/cosmos_manifest_example.py @@ -6,8 +6,8 @@ from datetime import datetime from pathlib import Path -from cosmos import DbtDag, ProjectConfig, ProfileConfig, RenderConfig, LoadMode, ExecutionConfig -from cosmos.profiles import PostgresUserPasswordProfileMapping, DbtProfileConfigVars +from cosmos import DbtDag, ExecutionConfig, LoadMode, ProfileConfig, ProjectConfig, RenderConfig +from cosmos.profiles import DbtProfileConfigVars, PostgresUserPasswordProfileMapping DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" DBT_ROOT_PATH = Path(os.getenv("DBT_ROOT_PATH", DEFAULT_DBT_ROOT_PATH)) diff --git a/dev/dags/cosmos_profile_mapping.py b/dev/dags/cosmos_profile_mapping.py index 48040126e..6570467e4 100644 --- a/dev/dags/cosmos_profile_mapping.py +++ b/dev/dags/cosmos_profile_mapping.py @@ -11,7 +11,7 @@ from airflow.decorators import dag from airflow.operators.empty import EmptyOperator -from cosmos import DbtTaskGroup, ProjectConfig, ProfileConfig +from cosmos import DbtTaskGroup, ProfileConfig, ProjectConfig from cosmos.profiles import get_automatic_profile_mapping DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" diff --git a/dev/dags/dbt/simple/convert_csv_to_db.py b/dev/dags/dbt/simple/convert_csv_to_db.py index 7424db32d..ef0ca301e 100644 --- a/dev/dags/dbt/simple/convert_csv_to_db.py +++ b/dev/dags/dbt/simple/convert_csv_to_db.py @@ -1,6 +1,7 @@ -import pandas as pd import sqlite3 +import pandas as pd + df = pd.read_csv("imdb.csv") conn = sqlite3.connect("imdb.db") df.to_sql("movies_ratings", conn, if_exists="replace", index=False) diff --git a/dev/dags/dbt_docs.py b/dev/dags/dbt_docs.py index edf89bdab..924928a80 100644 --- a/dev/dags/dbt_docs.py +++ b/dev/dags/dbt_docs.py @@ -11,16 +11,16 @@ from pathlib import Path from airflow import DAG -from airflow.hooks.base import BaseHook -from airflow.exceptions import AirflowNotFoundException from airflow.decorators import task +from airflow.exceptions import AirflowNotFoundException +from airflow.hooks.base import BaseHook from pendulum import datetime from cosmos import ProfileConfig from cosmos.operators import ( DbtDocsAzureStorageOperator, - DbtDocsS3Operator, DbtDocsGCSOperator, + DbtDocsS3Operator, ) from cosmos.profiles import PostgresUserPasswordProfileMapping diff --git a/dev/dags/example_cosmos_python_models.py b/dev/dags/example_cosmos_python_models.py index 3aa313913..90c8cc0ca 100644 --- a/dev/dags/example_cosmos_python_models.py +++ b/dev/dags/example_cosmos_python_models.py @@ -17,7 +17,7 @@ from datetime import datetime from pathlib import Path -from cosmos import DbtDag, ProjectConfig, ProfileConfig +from cosmos import DbtDag, ProfileConfig, ProjectConfig from cosmos.profiles import DatabricksTokenProfileMapping DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" diff --git a/dev/dags/example_cosmos_sources.py b/dev/dags/example_cosmos_sources.py index 24359ee3b..1a85b6d9f 100644 --- a/dev/dags/example_cosmos_sources.py +++ b/dev/dags/example_cosmos_sources.py @@ -16,11 +16,11 @@ from datetime import datetime from pathlib import Path -from airflow.operators.dummy import DummyOperator from airflow.models.dag import DAG +from airflow.operators.dummy import DummyOperator from airflow.utils.task_group import TaskGroup -from cosmos import DbtDag, ProjectConfig, ProfileConfig, RenderConfig +from cosmos import DbtDag, ProfileConfig, ProjectConfig, RenderConfig from cosmos.constants import DbtResourceType from cosmos.dbt.graph import DbtNode diff --git a/dev/dags/example_model_version.py b/dev/dags/example_model_version.py index 78f38647d..909c7494a 100644 --- a/dev/dags/example_model_version.py +++ b/dev/dags/example_model_version.py @@ -6,7 +6,7 @@ from datetime import datetime from pathlib import Path -from cosmos import DbtDag, ProjectConfig, ProfileConfig +from cosmos import DbtDag, ProfileConfig, ProjectConfig from cosmos.profiles import PostgresUserPasswordProfileMapping DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" diff --git a/dev/dags/example_virtualenv.py b/dev/dags/example_virtualenv.py index d84646cec..55ecf0a66 100644 --- a/dev/dags/example_virtualenv.py +++ b/dev/dags/example_virtualenv.py @@ -6,7 +6,7 @@ from datetime import datetime from pathlib import Path -from cosmos import DbtDag, ExecutionMode, ExecutionConfig, ProjectConfig, ProfileConfig +from cosmos import DbtDag, ExecutionConfig, ExecutionMode, ProfileConfig, ProjectConfig from cosmos.profiles import PostgresUserPasswordProfileMapping DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" diff --git a/dev/dags/performance_dag.py b/dev/dags/performance_dag.py index fec5175c8..1c8d639e0 100644 --- a/dev/dags/performance_dag.py +++ b/dev/dags/performance_dag.py @@ -2,14 +2,13 @@ An airflow DAG that uses Cosmos to render a dbt project for performance testing. """ -from datetime import datetime import os +from datetime import datetime from pathlib import Path -from cosmos import DbtDag, ProjectConfig, ProfileConfig, RenderConfig +from cosmos import DbtDag, ProfileConfig, ProjectConfig, RenderConfig from cosmos.profiles import PostgresUserPasswordProfileMapping - DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" DBT_ROOT_PATH = Path(os.getenv("DBT_ROOT_PATH", DEFAULT_DBT_ROOT_PATH)) diff --git a/dev/dags/user_defined_profile.py b/dev/dags/user_defined_profile.py index 22402bc53..a88354b04 100644 --- a/dev/dags/user_defined_profile.py +++ b/dev/dags/user_defined_profile.py @@ -9,7 +9,7 @@ from airflow.decorators import dag from airflow.operators.empty import EmptyOperator -from cosmos import DbtTaskGroup, ProjectConfig, ProfileConfig, RenderConfig, LoadMode +from cosmos import DbtTaskGroup, LoadMode, ProfileConfig, ProjectConfig, RenderConfig DEFAULT_DBT_ROOT_PATH = Path(__file__).parent / "dbt" DBT_ROOT_PATH = Path(os.getenv("DBT_ROOT_PATH", DEFAULT_DBT_ROOT_PATH)) diff --git a/docs/generate_mappings.py b/docs/generate_mappings.py index 273187ba5..8955efc73 100644 --- a/docs/generate_mappings.py +++ b/docs/generate_mappings.py @@ -9,7 +9,8 @@ from typing import Type from jinja2 import Environment, FileSystemLoader -from cosmos.profiles import profile_mappings, BaseProfileMapping + +from cosmos.profiles import BaseProfileMapping, profile_mappings @dataclass diff --git a/pyproject.toml b/pyproject.toml index 8b7d57aa5..0c26fa799 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -197,7 +197,7 @@ no_warn_unused_ignores = true [tool.ruff] line-length = 120 [tool.ruff.lint] -select = ["C901"] +select = ["C901", "I"] [tool.ruff.lint.mccabe] max-complexity = 8 diff --git a/tests/airflow/test_graph.py b/tests/airflow/test_graph.py index 16a2c7e07..1e5648306 100644 --- a/tests/airflow/test_graph.py +++ b/tests/airflow/test_graph.py @@ -9,13 +9,13 @@ from packaging import version from cosmos.airflow.graph import ( + _snake_case_to_camelcase, build_airflow_graph, calculate_leaves, calculate_operator_class, create_task_metadata, create_test_task_metadata, generate_task_or_group, - _snake_case_to_camelcase, ) from cosmos.config import ProfileConfig, RenderConfig from cosmos.constants import ( diff --git a/tests/dbt/parser/test_output.py b/tests/dbt/parser/test_output.py index 9fae4d3b3..4fe11669f 100644 --- a/tests/dbt/parser/test_output.py +++ b/tests/dbt/parser/test_output.py @@ -1,13 +1,14 @@ -import pytest import logging from unittest.mock import MagicMock + +import pytest from airflow.hooks.subprocess import SubprocessResult from cosmos.dbt.parser.output import ( extract_dbt_runner_issues, extract_log_issues, - parse_number_of_warnings_subprocess, parse_number_of_warnings_dbt_runner, + parse_number_of_warnings_subprocess, ) diff --git a/tests/dbt/test_graph.py b/tests/dbt/test_graph.py index 94951d7bd..4a0d4c98e 100644 --- a/tests/dbt/test_graph.py +++ b/tests/dbt/test_graph.py @@ -1,12 +1,13 @@ import shutil import tempfile from pathlib import Path -from unittest.mock import patch, MagicMock -import yaml +from subprocess import PIPE, Popen +from unittest.mock import MagicMock, patch import pytest +import yaml -from cosmos.config import ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig, CosmosConfigException +from cosmos.config import CosmosConfigException, ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig from cosmos.constants import DbtResourceType, ExecutionMode from cosmos.dbt.graph import ( CosmosLoadDbtException, @@ -17,7 +18,6 @@ run_command, ) from cosmos.profiles import PostgresUserPasswordProfileMapping -from subprocess import Popen, PIPE DBT_PROJECTS_ROOT_DIR = Path(__file__).parent.parent.parent / "dev/dags/dbt" DBT_PROJECT_NAME = "jaffle_shop" diff --git a/tests/dbt/test_project.py b/tests/dbt/test_project.py index a3cd30819..6f9e2cb84 100644 --- a/tests/dbt/test_project.py +++ b/tests/dbt/test_project.py @@ -4,7 +4,7 @@ import pytest -from cosmos.dbt.project import create_symlinks, copy_msgpack_for_partial_parse, environ, change_working_directory +from cosmos.dbt.project import change_working_directory, copy_msgpack_for_partial_parse, create_symlinks, environ DBT_PROJECTS_ROOT_DIR = Path(__file__).parent.parent.parent / "dev/dags/dbt" diff --git a/tests/dbt/test_selector.py b/tests/dbt/test_selector.py index ab7842783..bfb6d7d4e 100644 --- a/tests/dbt/test_selector.py +++ b/tests/dbt/test_selector.py @@ -2,10 +2,9 @@ import pytest -from cosmos.dbt.selector import SelectorConfig from cosmos.constants import DbtResourceType from cosmos.dbt.graph import DbtNode -from cosmos.dbt.selector import select_nodes +from cosmos.dbt.selector import SelectorConfig, select_nodes from cosmos.exceptions import CosmosValueError SAMPLE_PROJ_PATH = Path("/home/user/path/dbt-proj/") diff --git a/tests/hooks/test_subprocess.py b/tests/hooks/test_subprocess.py index 601d37b00..e8b16d387 100644 --- a/tests/hooks/test_subprocess.py +++ b/tests/hooks/test_subprocess.py @@ -1,7 +1,7 @@ +import signal from pathlib import Path from tempfile import TemporaryDirectory from unittest.mock import MagicMock, patch -import signal import pytest diff --git a/tests/operators/test_azure_container_instance.py b/tests/operators/test_azure_container_instance.py index da7720958..01fa3e20e 100644 --- a/tests/operators/test_azure_container_instance.py +++ b/tests/operators/test_azure_container_instance.py @@ -8,8 +8,8 @@ DbtAzureContainerInstanceBaseOperator, DbtLSAzureContainerInstanceOperator, DbtRunAzureContainerInstanceOperator, - DbtTestAzureContainerInstanceOperator, DbtSeedAzureContainerInstanceOperator, + DbtTestAzureContainerInstanceOperator, ) diff --git a/tests/operators/test_base.py b/tests/operators/test_base.py index 4ac78fc5a..edaf8a845 100644 --- a/tests/operators/test_base.py +++ b/tests/operators/test_base.py @@ -1,15 +1,16 @@ -import pytest from unittest.mock import patch +import pytest + from cosmos.operators.base import ( AbstractDbtBaseOperator, DbtBuildMixin, DbtLSMixin, - DbtSeedMixin, + DbtRunMixin, DbtRunOperationMixin, - DbtTestMixin, + DbtSeedMixin, DbtSnapshotMixin, - DbtRunMixin, + DbtTestMixin, ) diff --git a/tests/operators/test_docker.py b/tests/operators/test_docker.py index ef26fbaff..ad3ec5485 100644 --- a/tests/operators/test_docker.py +++ b/tests/operators/test_docker.py @@ -1,7 +1,7 @@ from pathlib import Path from unittest.mock import MagicMock, patch -import pytest +import pytest from airflow.utils.context import Context from pendulum import datetime diff --git a/tests/operators/test_kubernetes.py b/tests/operators/test_kubernetes.py index 19a6c7aeb..75739111f 100644 --- a/tests/operators/test_kubernetes.py +++ b/tests/operators/test_kubernetes.py @@ -2,6 +2,8 @@ from unittest.mock import MagicMock, patch import pytest +from airflow.models import TaskInstance +from airflow.utils.context import Context, context_merge from pendulum import datetime from cosmos.operators.kubernetes import ( @@ -12,9 +14,6 @@ DbtTestKubernetesOperator, ) -from airflow.utils.context import Context, context_merge -from airflow.models import TaskInstance - try: from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction diff --git a/tests/operators/test_local.py b/tests/operators/test_local.py index 9a938bca9..05cca4bcc 100644 --- a/tests/operators/test_local.py +++ b/tests/operators/test_local.py @@ -1,10 +1,10 @@ import logging import os -import sys import shutil +import sys import tempfile from pathlib import Path -from unittest.mock import MagicMock, patch, call +from unittest.mock import MagicMock, call, patch import pytest from airflow import DAG @@ -16,28 +16,28 @@ from pendulum import datetime from cosmos.config import ProfileConfig +from cosmos.constants import InvocationMode +from cosmos.dbt.parser.output import ( + extract_dbt_runner_issues, + parse_number_of_warnings_dbt_runner, + parse_number_of_warnings_subprocess, +) from cosmos.operators.local import ( - DbtLocalBaseOperator, - DbtLSLocalOperator, - DbtSnapshotLocalOperator, - DbtRunLocalOperator, - DbtTestLocalOperator, DbtBuildLocalOperator, - DbtDocsLocalOperator, - DbtDocsS3LocalOperator, DbtDocsAzureStorageLocalOperator, DbtDocsGCSLocalOperator, - DbtSeedLocalOperator, + DbtDocsLocalOperator, + DbtDocsS3LocalOperator, + DbtLocalBaseOperator, + DbtLSLocalOperator, + DbtRunLocalOperator, DbtRunOperationLocalOperator, + DbtSeedLocalOperator, + DbtSnapshotLocalOperator, + DbtTestLocalOperator, ) from cosmos.profiles import PostgresUserPasswordProfileMapping -from cosmos.constants import InvocationMode from tests.utils import test_dag as run_test_dag -from cosmos.dbt.parser.output import ( - extract_dbt_runner_issues, - parse_number_of_warnings_subprocess, - parse_number_of_warnings_dbt_runner, -) DBT_PROJ_DIR = Path(__file__).parent.parent.parent / "dev/dags/dbt/jaffle_shop" MINI_DBT_PROJ_DIR = Path(__file__).parent.parent / "sample/mini" diff --git a/tests/operators/test_virtualenv.py b/tests/operators/test_virtualenv.py index 036f162de..1610d2707 100644 --- a/tests/operators/test_virtualenv.py +++ b/tests/operators/test_virtualenv.py @@ -1,13 +1,11 @@ -from unittest.mock import patch, MagicMock - -from cosmos.operators.virtualenv import DbtVirtualenvBaseOperator +from unittest.mock import MagicMock, patch from airflow.models.connection import Connection from cosmos.config import ProfileConfig - -from cosmos.profiles import PostgresUserPasswordProfileMapping from cosmos.constants import InvocationMode +from cosmos.operators.virtualenv import DbtVirtualenvBaseOperator +from cosmos.profiles import PostgresUserPasswordProfileMapping profile_config = ProfileConfig( profile_name="default", diff --git a/tests/perf/test_performance.py b/tests/perf/test_performance.py index 81b08d8bd..995c33a74 100644 --- a/tests/perf/test_performance.py +++ b/tests/perf/test_performance.py @@ -1,9 +1,9 @@ from __future__ import annotations -import time import os -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path from typing import Generator try: diff --git a/tests/plugin/test_plugin.py b/tests/plugin/test_plugin.py index df33ae13a..8d0e3742f 100644 --- a/tests/plugin/test_plugin.py +++ b/tests/plugin/test_plugin.py @@ -7,15 +7,15 @@ try: import flask # noqa: F401 except ImportError: - import markupsafe import jinja2 + import markupsafe jinja2.Markup = markupsafe.Markup jinja2.escape = markupsafe.escape -from unittest.mock import mock_open, patch, MagicMock, PropertyMock - import sys +from unittest.mock import MagicMock, PropertyMock, mock_open, patch + import pytest from airflow.configuration import conf from airflow.exceptions import AirflowConfigException @@ -25,18 +25,16 @@ from flask.testing import FlaskClient import cosmos.plugin - from cosmos.plugin import ( dbt_docs_view, iframe_script, - open_gcs_file, open_azure_file, + open_file, + open_gcs_file, open_http_file, open_s3_file, - open_file, ) - original_conf_get = conf.get diff --git a/tests/profiles/athena/test_athena_access_key.py b/tests/profiles/athena/test_athena_access_key.py index c224a9d4b..71ba1eb05 100644 --- a/tests/profiles/athena/test_athena_access_key.py +++ b/tests/profiles/athena/test_athena_access_key.py @@ -1,9 +1,10 @@ "Tests for the Athena profile." import json -from collections import namedtuple import sys +from collections import namedtuple from unittest.mock import MagicMock, patch + import pytest from airflow.models.connection import Connection diff --git a/tests/profiles/bigquery/test_bq_service_account_keyfile_dict.py b/tests/profiles/bigquery/test_bq_service_account_keyfile_dict.py index 0647e6800..00cf070c3 100644 --- a/tests/profiles/bigquery/test_bq_service_account_keyfile_dict.py +++ b/tests/profiles/bigquery/test_bq_service_account_keyfile_dict.py @@ -3,8 +3,8 @@ import pytest from airflow.models.connection import Connection -from cosmos.exceptions import CosmosValueError +from cosmos.exceptions import CosmosValueError from cosmos.profiles import get_automatic_profile_mapping from cosmos.profiles.bigquery.service_account_keyfile_dict import GoogleCloudServiceAccountDictProfileMapping diff --git a/tests/profiles/test_base_profile.py b/tests/profiles/test_base_profile.py index b80912bcd..7fdbdb886 100644 --- a/tests/profiles/test_base_profile.py +++ b/tests/profiles/test_base_profile.py @@ -1,12 +1,13 @@ from __future__ import annotations + from typing import Any import pytest import yaml from pydantic.error_wrappers import ValidationError -from cosmos.profiles.base import BaseProfileMapping, DbtProfileConfigVars from cosmos.exceptions import CosmosValueError +from cosmos.profiles.base import BaseProfileMapping, DbtProfileConfigVars class TestProfileMapping(BaseProfileMapping): diff --git a/tests/profiles/trino/test_trino_base.py b/tests/profiles/trino/test_trino_base.py index 29e402691..31f1a3166 100644 --- a/tests/profiles/trino/test_trino_base.py +++ b/tests/profiles/trino/test_trino_base.py @@ -1,7 +1,7 @@ "Tests for the Trino profile." -from unittest.mock import patch import json +from unittest.mock import patch from airflow.models.connection import Connection diff --git a/tests/profiles/trino/test_trino_jwt.py b/tests/profiles/trino/test_trino_jwt.py index d62fefc47..36d2d2153 100644 --- a/tests/profiles/trino/test_trino_jwt.py +++ b/tests/profiles/trino/test_trino_jwt.py @@ -6,8 +6,7 @@ import pytest from airflow.models.connection import Connection -from cosmos.profiles import get_automatic_profile_mapping -from cosmos.profiles import TrinoJWTProfileMapping +from cosmos.profiles import TrinoJWTProfileMapping, get_automatic_profile_mapping @pytest.fixture() diff --git a/tests/profiles/trino/test_trino_ldap.py b/tests/profiles/trino/test_trino_ldap.py index a959e9e68..f8adddf93 100644 --- a/tests/profiles/trino/test_trino_ldap.py +++ b/tests/profiles/trino/test_trino_ldap.py @@ -5,8 +5,7 @@ import pytest from airflow.models.connection import Connection -from cosmos.profiles import get_automatic_profile_mapping -from cosmos.profiles import TrinoLDAPProfileMapping +from cosmos.profiles import TrinoLDAPProfileMapping, get_automatic_profile_mapping @pytest.fixture() diff --git a/tests/test_config.py b/tests/test_config.py index b93ad2627..acca546be 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,14 +1,13 @@ +from contextlib import nullcontext as does_not_raise from pathlib import Path from unittest.mock import patch -from cosmos.profiles.postgres.user_pass import PostgresUserPasswordProfileMapping -from contextlib import nullcontext as does_not_raise import pytest +from cosmos.config import CosmosConfigException, ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig from cosmos.constants import ExecutionMode, InvocationMode -from cosmos.config import ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig, CosmosConfigException from cosmos.exceptions import CosmosValueError - +from cosmos.profiles.postgres.user_pass import PostgresUserPasswordProfileMapping DBT_PROJECTS_ROOT_DIR = Path(__file__).parent / "sample/" SAMPLE_PROFILE_YML = Path(__file__).parent / "sample/profiles.yml" diff --git a/tests/test_converter.py b/tests/test_converter.py index e66af468f..f5dfe68b9 100644 --- a/tests/test_converter.py +++ b/tests/test_converter.py @@ -1,17 +1,16 @@ from datetime import datetime from pathlib import Path -from unittest.mock import patch, MagicMock -from cosmos.profiles.postgres import PostgresUserPasswordProfileMapping +from unittest.mock import MagicMock, patch import pytest from airflow.models import DAG +from cosmos.config import CosmosConfigException, ExecutionConfig, ProfileConfig, ProjectConfig, RenderConfig +from cosmos.constants import DbtResourceType, ExecutionMode, InvocationMode, LoadMode from cosmos.converter import DbtToAirflowConverter, validate_arguments, validate_initial_user_config -from cosmos.constants import DbtResourceType, ExecutionMode, LoadMode, InvocationMode -from cosmos.config import ProjectConfig, ProfileConfig, ExecutionConfig, RenderConfig, CosmosConfigException from cosmos.dbt.graph import DbtNode from cosmos.exceptions import CosmosValueError - +from cosmos.profiles.postgres import PostgresUserPasswordProfileMapping SAMPLE_PROFILE_YML = Path(__file__).parent / "sample/profiles.yml" SAMPLE_DBT_PROJECT = Path(__file__).parent / "sample/" diff --git a/tests/test_example_dags_no_connections.py b/tests/test_example_dags_no_connections.py index ae7c354a1..f5b5c0845 100644 --- a/tests/test_example_dags_no_connections.py +++ b/tests/test_example_dags_no_connections.py @@ -13,7 +13,6 @@ from dbt.version import get_installed_version as get_dbt_version from packaging.version import Version - EXAMPLE_DAGS_DIR = Path(__file__).parent.parent / "dev/dags" AIRFLOW_IGNORE_FILE = EXAMPLE_DAGS_DIR / ".airflowignore" DBT_VERSION = Version(get_dbt_version().to_version_string()[1:]) diff --git a/tests/test_log.py b/tests/test_log.py index d145dbca4..c110c1918 100644 --- a/tests/test_log.py +++ b/tests/test_log.py @@ -1,8 +1,9 @@ import logging +from airflow.configuration import conf + from cosmos import get_provider_info from cosmos.log import get_logger -from airflow.configuration import conf def test_get_logger():