From c6b91b48ea32a8666bbbbafbdf64bdf2dcc64619 Mon Sep 17 00:00:00 2001 From: kenxu95 Date: Tue, 18 Apr 2023 11:58:55 -0400 Subject: [PATCH] lint --- examples/run_notebook.py | 3 ++- .../setup/flow_with_metrics_and_checks.py | 3 ++- integration_tests/backend/utils.py | 3 ++- .../no_concurrency/sdk/conftest.py | 2 +- .../no_concurrency/sdk/requirements_test.py | 2 +- .../sdk/aqueduct_tests/checks_test.py | 4 +-- .../sdk/aqueduct_tests/conftest.py | 1 - .../sdk/aqueduct_tests/data_validator.py | 2 +- .../aqueduct_tests/delete_workflow_test.py | 4 +-- .../sdk/aqueduct_tests/error_handling_test.py | 2 +- .../aqueduct_tests/external_compute_test.py | 2 +- .../sdk/aqueduct_tests/extract.py | 1 + .../sdk/aqueduct_tests/flow_test.py | 4 +-- .../sdk/aqueduct_tests/integration_test.py | 6 ++--- .../sdk/aqueduct_tests/lazy_execution_test.py | 4 +-- .../sdk/aqueduct_tests/metrics_test.py | 2 +- .../aqueduct_tests/multiple_output_test.py | 2 +- .../sdk/aqueduct_tests/naming_test.py | 2 +- .../aqueduct_tests/operator_hierarchy_test.py | 2 +- .../sdk/aqueduct_tests/operator_test.py | 4 +-- .../sdk/aqueduct_tests/param_test.py | 8 +++--- .../sdk/aqueduct_tests/preview_test.py | 2 +- .../sdk/aqueduct_tests/resources_test.py | 4 +-- integration_tests/sdk/aqueduct_tests/save.py | 3 ++- .../sdk/aqueduct_tests/spark_test.py | 4 +-- .../aqueduct_tests/type_enforcement_test.py | 4 +-- integration_tests/sdk/conftest.py | 4 +-- .../sdk/data_integration_tests/athena_test.py | 4 +-- .../sdk/data_integration_tests/conftest.py | 1 + .../data_integration_tests/flow_manager.py | 3 ++- .../mongo_db_data_validator.py | 2 +- .../data_integration_tests/mongo_db_test.py | 4 +-- .../relational_data_validator.py | 4 +-- .../data_integration_tests/relational_test.py | 4 +-- .../s3_data_validator.py | 4 +-- .../sdk/data_integration_tests/s3_test.py | 6 ++--- integration_tests/sdk/setup_integration.py | 4 +-- integration_tests/sdk/shared/flow_helpers.py | 5 ++-- integration_tests/sdk/shared/utils.py | 5 ++-- integration_tests/sdk/shared/validation.py | 3 ++- integration_tests/sdk/shared/validator.py | 4 +-- manual_qa_tests/initialize.py | 2 +- manual_qa_tests/notification.py | 3 ++- .../succeed_march_madness_dag_layout_test.py | 3 ++- sdk/aqueduct/artifacts/bool_artifact.py | 1 - sdk/aqueduct/artifacts/generic_artifact.py | 3 ++- sdk/aqueduct/artifacts/numeric_artifact.py | 4 +-- sdk/aqueduct/artifacts/preview.py | 3 ++- sdk/aqueduct/artifacts/system_metric.py | 3 ++- sdk/aqueduct/artifacts/table_artifact.py | 6 ++--- sdk/aqueduct/backend/api_client.py | 3 +-- sdk/aqueduct/backend/response_helpers.py | 3 +-- sdk/aqueduct/client.py | 4 +-- sdk/aqueduct/decorator.py | 4 +-- sdk/aqueduct/error.py | 4 ++- sdk/aqueduct/flow.py | 3 ++- sdk/aqueduct/flow_run.py | 3 ++- sdk/aqueduct/github.py | 3 ++- sdk/aqueduct/integrations/connect_config.py | 3 +-- .../integrations/dynamic_k8s_integration.py | 6 ++--- .../integrations/google_sheets_integration.py | 3 +-- .../integrations/mongodb_integration.py | 6 ++--- sdk/aqueduct/integrations/s3_integration.py | 6 ++--- .../integrations/salesforce_integration.py | 3 +-- sdk/aqueduct/integrations/sql_integration.py | 6 ++--- sdk/aqueduct/integrations/validation.py | 12 ++++++--- sdk/aqueduct/models/artifact.py | 3 +-- sdk/aqueduct/models/config.py | 3 +-- sdk/aqueduct/models/dag.py | 3 +-- sdk/aqueduct/models/integration.py | 27 +------------------ sdk/aqueduct/models/operators.py | 3 +-- sdk/aqueduct/models/response_models.py | 3 +-- sdk/aqueduct/tests/decorator_test.py | 1 - ...ecorators_with_without_parentheses_test.py | 4 +-- sdk/aqueduct/tests/flow_test.py | 1 - sdk/aqueduct/tests/metric_test.py | 3 ++- sdk/aqueduct/tests/serialization_test.py | 3 +-- sdk/aqueduct/tests/utils.py | 4 +-- sdk/aqueduct/utils/function_packaging.py | 1 - sdk/aqueduct/utils/integration_validation.py | 15 ++++++----- sdk/aqueduct/utils/serialization.py | 7 +++-- sdk/aqueduct/utils/type_inference.py | 3 +-- sdk/aqueduct/utils/utils.py | 6 ++--- .../artifact_migration_000016/execute.py | 3 +-- .../execute.py | 3 +-- .../operators/connectors/data/extract.py | 3 +-- .../operators/connectors/data/s3.py | 3 +-- .../connectors/data/s3_serialization.py | 7 +++-- .../operators/function_executor/execute.py | 5 ++-- .../operators/param_executor/execute.py | 5 ++-- .../operators/spark/execute_function.py | 3 +-- .../operators/spark/utils.py | 13 +++++---- .../operators/utils/utils.py | 13 +++++---- 93 files changed, 179 insertions(+), 206 deletions(-) diff --git a/examples/run_notebook.py b/examples/run_notebook.py index 3da69897e..f4ca6d7f0 100644 --- a/examples/run_notebook.py +++ b/examples/run_notebook.py @@ -7,9 +7,10 @@ from pathlib import Path from typing import List -from aqueduct import Client, get_apikey from aqueduct.error import InvalidUserArgumentException +from aqueduct import Client, get_apikey + """ See README.md for details about this script. """ diff --git a/integration_tests/backend/setup/flow_with_metrics_and_checks.py b/integration_tests/backend/setup/flow_with_metrics_and_checks.py index 91b2fc7dd..80355c625 100644 --- a/integration_tests/backend/setup/flow_with_metrics_and_checks.py +++ b/integration_tests/backend/setup/flow_with_metrics_and_checks.py @@ -3,9 +3,10 @@ # with a check to enforce the row-count is larger than 0. # This workflow is published twice. ### -import aqueduct from aqueduct.constants.enums import CheckSeverity +import aqueduct + def setup_flow_with_metrics_and_checks( client: aqueduct.Client, diff --git a/integration_tests/backend/utils.py b/integration_tests/backend/utils.py index f06116846..8a2105a99 100644 --- a/integration_tests/backend/utils.py +++ b/integration_tests/backend/utils.py @@ -1,9 +1,10 @@ import time from typing import Optional -import aqueduct from aqueduct.constants.enums import ExecutionStatus +import aqueduct + def wait_for_flow_runs( client: aqueduct.Client, flow_id: str, num_runs: int = 1, expect_success: Optional[bool] = None diff --git a/integration_tests/no_concurrency/sdk/conftest.py b/integration_tests/no_concurrency/sdk/conftest.py index bff6d02fc..5a1cff3e0 100644 --- a/integration_tests/no_concurrency/sdk/conftest.py +++ b/integration_tests/no_concurrency/sdk/conftest.py @@ -1,10 +1,10 @@ import os import pytest +from aqueduct.models.dag import DAG, Metadata import aqueduct from aqueduct import globals -from aqueduct.models.dag import DAG, Metadata API_KEY_ENV_NAME = "API_KEY" SERVER_ADDR_ENV_NAME = "SERVER_ADDRESS" diff --git a/integration_tests/no_concurrency/sdk/requirements_test.py b/integration_tests/no_concurrency/sdk/requirements_test.py index 72813e446..e31eed70d 100644 --- a/integration_tests/no_concurrency/sdk/requirements_test.py +++ b/integration_tests/no_concurrency/sdk/requirements_test.py @@ -2,11 +2,11 @@ import pandas as pd import pytest +from aqueduct.error import AqueductError, InvalidUserArgumentException from transformers_model.model import sentiment_prediction_using_transformers from utils import SENTIMENT_SQL_QUERY, get_integration_name from aqueduct import op -from aqueduct.error import AqueductError, InvalidUserArgumentException INVALID_REQUIREMENTS_PATH = "~/random.txt" VALID_REQUIREMENTS_PATH = "transformers_model/requirements.txt" diff --git a/integration_tests/sdk/aqueduct_tests/checks_test.py b/integration_tests/sdk/aqueduct_tests/checks_test.py index b47dbeee8..a09f09dee 100644 --- a/integration_tests/sdk/aqueduct_tests/checks_test.py +++ b/integration_tests/sdk/aqueduct_tests/checks_test.py @@ -1,7 +1,5 @@ import pandas as pd import pytest - -from aqueduct import check from aqueduct.constants.enums import CheckSeverity from aqueduct.error import ( AqueductError, @@ -10,6 +8,8 @@ InvalidUserActionException, ) +from aqueduct import check + from ..shared.data_objects import DataObject from ..shared.flow_helpers import publish_flow_test from .extract import extract diff --git a/integration_tests/sdk/aqueduct_tests/conftest.py b/integration_tests/sdk/aqueduct_tests/conftest.py index 941432c0b..93e704454 100644 --- a/integration_tests/sdk/aqueduct_tests/conftest.py +++ b/integration_tests/sdk/aqueduct_tests/conftest.py @@ -1,5 +1,4 @@ import pytest - from aqueduct.constants.enums import ServiceType from .data_validator import DataValidator diff --git a/integration_tests/sdk/aqueduct_tests/data_validator.py b/integration_tests/sdk/aqueduct_tests/data_validator.py index 9ff2ac1b7..d9a174bb7 100644 --- a/integration_tests/sdk/aqueduct_tests/data_validator.py +++ b/integration_tests/sdk/aqueduct_tests/data_validator.py @@ -2,9 +2,9 @@ from typing import Any import pandas as pd +from aqueduct.models.integration import Integration from aqueduct import Client, Flow -from aqueduct.models.integration import Integration from ..shared.globals import artifact_id_to_saved_identifier from ..shared.validation import fetch_and_validate_saved_object_identifier diff --git a/integration_tests/sdk/aqueduct_tests/delete_workflow_test.py b/integration_tests/sdk/aqueduct_tests/delete_workflow_test.py index e52f0339f..330489063 100644 --- a/integration_tests/sdk/aqueduct_tests/delete_workflow_test.py +++ b/integration_tests/sdk/aqueduct_tests/delete_workflow_test.py @@ -1,10 +1,10 @@ import pandas as pd import pytest +from aqueduct.constants.enums import LoadUpdateMode +from aqueduct.error import InvalidRequestError, InvalidUserArgumentException import aqueduct from aqueduct import op -from aqueduct.constants.enums import LoadUpdateMode -from aqueduct.error import InvalidRequestError, InvalidUserArgumentException from ..shared.data_objects import DataObject from ..shared.flow_helpers import publish_flow_test diff --git a/integration_tests/sdk/aqueduct_tests/error_handling_test.py b/integration_tests/sdk/aqueduct_tests/error_handling_test.py index d96bf408c..f52a6699e 100644 --- a/integration_tests/sdk/aqueduct_tests/error_handling_test.py +++ b/integration_tests/sdk/aqueduct_tests/error_handling_test.py @@ -1,8 +1,8 @@ import pytest +from aqueduct.error import AqueductError, InvalidUserArgumentException import aqueduct from aqueduct import op -from aqueduct.error import AqueductError, InvalidUserArgumentException from ..shared.data_objects import DataObject from .extract import extract diff --git a/integration_tests/sdk/aqueduct_tests/external_compute_test.py b/integration_tests/sdk/aqueduct_tests/external_compute_test.py index ed04ce1c5..e18f6e53c 100644 --- a/integration_tests/sdk/aqueduct_tests/external_compute_test.py +++ b/integration_tests/sdk/aqueduct_tests/external_compute_test.py @@ -1,9 +1,9 @@ import platform import pytest +from aqueduct.error import AqueductError from aqueduct import check, global_config, op -from aqueduct.error import AqueductError from ..shared.data_objects import DataObject from ..shared.flow_helpers import publish_flow_test diff --git a/integration_tests/sdk/aqueduct_tests/extract.py b/integration_tests/sdk/aqueduct_tests/extract.py index 779acaeeb..3686ea7f0 100644 --- a/integration_tests/sdk/aqueduct_tests/extract.py +++ b/integration_tests/sdk/aqueduct_tests/extract.py @@ -4,6 +4,7 @@ from aqueduct.constants.enums import ArtifactType from aqueduct.integrations.s3_integration import S3Integration from aqueduct.integrations.sql_integration import RelationalDBIntegration + from sdk.shared.data_objects import DataObject diff --git a/integration_tests/sdk/aqueduct_tests/flow_test.py b/integration_tests/sdk/aqueduct_tests/flow_test.py index a7d598b63..8d2b43a04 100644 --- a/integration_tests/sdk/aqueduct_tests/flow_test.py +++ b/integration_tests/sdk/aqueduct_tests/flow_test.py @@ -3,11 +3,11 @@ import pandas as pd import pytest +from aqueduct.constants.enums import ExecutionStatus +from aqueduct.error import InvalidRequestError, InvalidUserArgumentException import aqueduct from aqueduct import check, metric, op -from aqueduct.constants.enums import ExecutionStatus -from aqueduct.error import InvalidRequestError, InvalidUserArgumentException from ..shared.data_objects import DataObject from ..shared.flow_helpers import publish_flow_test, trigger_flow_test, wait_for_flow_runs diff --git a/integration_tests/sdk/aqueduct_tests/integration_test.py b/integration_tests/sdk/aqueduct_tests/integration_test.py index 07e82a60c..431e384d7 100644 --- a/integration_tests/sdk/aqueduct_tests/integration_test.py +++ b/integration_tests/sdk/aqueduct_tests/integration_test.py @@ -1,7 +1,4 @@ import pytest -from pydantic import ValidationError - -from aqueduct import global_config from aqueduct.constants.enums import ServiceType from aqueduct.error import ( AqueductError, @@ -11,6 +8,9 @@ InvalidUserArgumentException, ) from aqueduct.integrations.connect_config import K8sConfig +from pydantic import ValidationError + +from aqueduct import global_config from ..shared.data_objects import DataObject from .extract import extract diff --git a/integration_tests/sdk/aqueduct_tests/lazy_execution_test.py b/integration_tests/sdk/aqueduct_tests/lazy_execution_test.py index 97eadc244..6cffd7bfc 100644 --- a/integration_tests/sdk/aqueduct_tests/lazy_execution_test.py +++ b/integration_tests/sdk/aqueduct_tests/lazy_execution_test.py @@ -2,14 +2,14 @@ import pandas as pd import pytest - -from aqueduct import check, global_config, metric, op from aqueduct.artifacts.bool_artifact import BoolArtifact from aqueduct.artifacts.generic_artifact import GenericArtifact from aqueduct.artifacts.numeric_artifact import NumericArtifact from aqueduct.constants.enums import ArtifactType from aqueduct.error import InvalidUserArgumentException +from aqueduct import check, global_config, metric, op + from ..shared.data_objects import DataObject from ..shared.flow_helpers import publish_flow_test from .extract import extract diff --git a/integration_tests/sdk/aqueduct_tests/metrics_test.py b/integration_tests/sdk/aqueduct_tests/metrics_test.py index 1c7e3d989..a134b5f2d 100644 --- a/integration_tests/sdk/aqueduct_tests/metrics_test.py +++ b/integration_tests/sdk/aqueduct_tests/metrics_test.py @@ -1,8 +1,8 @@ import pandas as pd import pytest +from aqueduct.error import AqueductError, ArtifactNotFoundException from aqueduct import metric -from aqueduct.error import AqueductError, ArtifactNotFoundException from ..shared.data_objects import DataObject from ..shared.flow_helpers import publish_flow_test diff --git a/integration_tests/sdk/aqueduct_tests/multiple_output_test.py b/integration_tests/sdk/aqueduct_tests/multiple_output_test.py index ce934192a..1c1d76153 100644 --- a/integration_tests/sdk/aqueduct_tests/multiple_output_test.py +++ b/integration_tests/sdk/aqueduct_tests/multiple_output_test.py @@ -1,7 +1,7 @@ import pytest +from aqueduct.error import AqueductError from aqueduct import op -from aqueduct.error import AqueductError from ..shared.flow_helpers import publish_flow_test diff --git a/integration_tests/sdk/aqueduct_tests/naming_test.py b/integration_tests/sdk/aqueduct_tests/naming_test.py index 0e32b3001..9847f8656 100644 --- a/integration_tests/sdk/aqueduct_tests/naming_test.py +++ b/integration_tests/sdk/aqueduct_tests/naming_test.py @@ -1,7 +1,7 @@ import pytest +from aqueduct.error import ArtifactNotFoundException, InvalidUserActionException from aqueduct import check, metric, op -from aqueduct.error import ArtifactNotFoundException, InvalidUserActionException from ..shared.data_objects import DataObject from ..shared.flow_helpers import publish_flow_test diff --git a/integration_tests/sdk/aqueduct_tests/operator_hierarchy_test.py b/integration_tests/sdk/aqueduct_tests/operator_hierarchy_test.py index 8ff68c5b7..da50a45f3 100644 --- a/integration_tests/sdk/aqueduct_tests/operator_hierarchy_test.py +++ b/integration_tests/sdk/aqueduct_tests/operator_hierarchy_test.py @@ -1,7 +1,7 @@ import pytest +from aqueduct.error import InvalidUserActionException from aqueduct import check, metric, op -from aqueduct.error import InvalidUserActionException from ..shared.data_objects import DataObject from .extract import extract diff --git a/integration_tests/sdk/aqueduct_tests/operator_test.py b/integration_tests/sdk/aqueduct_tests/operator_test.py index b394369b7..366b1bf33 100644 --- a/integration_tests/sdk/aqueduct_tests/operator_test.py +++ b/integration_tests/sdk/aqueduct_tests/operator_test.py @@ -1,9 +1,9 @@ import pandas as pd import pytest - -from aqueduct import op from aqueduct.decorator import to_operator from aqueduct.error import ArtifactNotFoundException + +from aqueduct import op from sdk.aqueduct_tests.test_function import dummy_sentiment_model_function from ..shared.data_objects import DataObject diff --git a/integration_tests/sdk/aqueduct_tests/param_test.py b/integration_tests/sdk/aqueduct_tests/param_test.py index 0510de622..2d39a0390 100644 --- a/integration_tests/sdk/aqueduct_tests/param_test.py +++ b/integration_tests/sdk/aqueduct_tests/param_test.py @@ -4,10 +4,6 @@ import pandas as pd import pytest -from pandas._testing import assert_frame_equal -from PIL import Image - -from aqueduct import metric, op from aqueduct.artifacts.generic_artifact import GenericArtifact from aqueduct.artifacts.numeric_artifact import NumericArtifact from aqueduct.constants.enums import ArtifactType, ExecutionStatus @@ -17,6 +13,10 @@ InvalidUserActionException, InvalidUserArgumentException, ) +from pandas._testing import assert_frame_equal +from PIL import Image + +from aqueduct import metric, op from ..shared.data_objects import DataObject from ..shared.flow_helpers import publish_flow_test, trigger_flow_test diff --git a/integration_tests/sdk/aqueduct_tests/preview_test.py b/integration_tests/sdk/aqueduct_tests/preview_test.py index af8366b83..28f0d3347 100644 --- a/integration_tests/sdk/aqueduct_tests/preview_test.py +++ b/integration_tests/sdk/aqueduct_tests/preview_test.py @@ -1,8 +1,8 @@ import pandas as pd import pytest +from aqueduct.error import AqueductError, InvalidDependencyFilePath, InvalidFunctionException from aqueduct import global_config, op -from aqueduct.error import AqueductError, InvalidDependencyFilePath, InvalidFunctionException from ..shared.data_objects import DataObject from .extract import extract diff --git a/integration_tests/sdk/aqueduct_tests/resources_test.py b/integration_tests/sdk/aqueduct_tests/resources_test.py index e6f27df5b..c59400e78 100644 --- a/integration_tests/sdk/aqueduct_tests/resources_test.py +++ b/integration_tests/sdk/aqueduct_tests/resources_test.py @@ -1,11 +1,11 @@ from os import cpu_count import pytest - -from aqueduct import global_config, op from aqueduct.constants.enums import ExecutionStatus, ServiceType from aqueduct.error import AqueductError, InvalidUserArgumentException +from aqueduct import global_config, op + from ..shared.flow_helpers import publish_flow_test diff --git a/integration_tests/sdk/aqueduct_tests/save.py b/integration_tests/sdk/aqueduct_tests/save.py index 2ddb06eef..18722fab2 100644 --- a/integration_tests/sdk/aqueduct_tests/save.py +++ b/integration_tests/sdk/aqueduct_tests/save.py @@ -1,12 +1,13 @@ from typing import Optional -from aqueduct import LoadUpdateMode from aqueduct.artifacts.table_artifact import TableArtifact from aqueduct.constants.enums import S3TableFormat from aqueduct.integrations.mongodb_integration import MongoDBIntegration from aqueduct.integrations.s3_integration import S3Integration from aqueduct.integrations.sql_integration import RelationalDBIntegration +from aqueduct import LoadUpdateMode + from ..shared.globals import artifact_id_to_saved_identifier from ..shared.naming import generate_table_name diff --git a/integration_tests/sdk/aqueduct_tests/spark_test.py b/integration_tests/sdk/aqueduct_tests/spark_test.py index a50598c79..a93dcb807 100644 --- a/integration_tests/sdk/aqueduct_tests/spark_test.py +++ b/integration_tests/sdk/aqueduct_tests/spark_test.py @@ -1,11 +1,11 @@ from os import cpu_count import pytest - -from aqueduct import global_config, op from aqueduct.constants.enums import ExecutionStatus, ServiceType from aqueduct.error import AqueductError, InvalidUserArgumentException +from aqueduct import global_config, op + from ..shared.data_objects import DataObject from ..shared.flow_helpers import publish_flow_test from .extract import extract diff --git a/integration_tests/sdk/aqueduct_tests/type_enforcement_test.py b/integration_tests/sdk/aqueduct_tests/type_enforcement_test.py index 9759a6b80..418b875f6 100644 --- a/integration_tests/sdk/aqueduct_tests/type_enforcement_test.py +++ b/integration_tests/sdk/aqueduct_tests/type_enforcement_test.py @@ -1,11 +1,11 @@ from typing import Union import pytest - -from aqueduct import op from aqueduct.constants.enums import ExecutionStatus from aqueduct.error import AqueductError +from aqueduct import op + from ..shared.flow_helpers import publish_flow_test, trigger_flow_test diff --git a/integration_tests/sdk/conftest.py b/integration_tests/sdk/conftest.py index c885a43e9..8855ec77d 100644 --- a/integration_tests/sdk/conftest.py +++ b/integration_tests/sdk/conftest.py @@ -1,10 +1,10 @@ from typing import Optional import pytest - -from aqueduct import Client, global_config, globals from aqueduct.constants.enums import ServiceType from aqueduct.models.dag import DAG, Metadata + +from aqueduct import Client, global_config, globals from sdk.setup_integration import ( get_aqueduct_config, get_artifact_store_name, diff --git a/integration_tests/sdk/data_integration_tests/athena_test.py b/integration_tests/sdk/data_integration_tests/athena_test.py index 40c0dd8fe..0f1be975c 100644 --- a/integration_tests/sdk/data_integration_tests/athena_test.py +++ b/integration_tests/sdk/data_integration_tests/athena_test.py @@ -1,10 +1,10 @@ import pandas as pd import pytest - -from aqueduct import LoadUpdateMode, metric, op from aqueduct.error import InvalidUserActionException, InvalidUserArgumentException from aqueduct.integrations.sql_integration import RelationalDBIntegration +from aqueduct import LoadUpdateMode, metric, op + from ..shared.naming import generate_table_name from .save import save diff --git a/integration_tests/sdk/data_integration_tests/conftest.py b/integration_tests/sdk/data_integration_tests/conftest.py index 157284ad7..fdfb5e4ca 100644 --- a/integration_tests/sdk/data_integration_tests/conftest.py +++ b/integration_tests/sdk/data_integration_tests/conftest.py @@ -5,6 +5,7 @@ # Maps the test files in this directory to the allowed data integrations for that file. # If a disallowed data integration is used, all tests in the file will be skipped. from aqueduct.constants.enums import ServiceType + from sdk.data_integration_tests.flow_manager import FlowManager allowed_data_integrations_by_file = { diff --git a/integration_tests/sdk/data_integration_tests/flow_manager.py b/integration_tests/sdk/data_integration_tests/flow_manager.py index 44e91d577..e5b61b902 100644 --- a/integration_tests/sdk/data_integration_tests/flow_manager.py +++ b/integration_tests/sdk/data_integration_tests/flow_manager.py @@ -1,8 +1,9 @@ from typing import Any, Callable, Dict, List, Optional, Union -from aqueduct import Client, Flow from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.constants.enums import ExecutionStatus + +from aqueduct import Client, Flow from sdk.shared.flow_helpers import publish_flow_test, trigger_flow_test diff --git a/integration_tests/sdk/data_integration_tests/mongo_db_data_validator.py b/integration_tests/sdk/data_integration_tests/mongo_db_data_validator.py index 1da2e71c9..37b7a03dd 100644 --- a/integration_tests/sdk/data_integration_tests/mongo_db_data_validator.py +++ b/integration_tests/sdk/data_integration_tests/mongo_db_data_validator.py @@ -2,9 +2,9 @@ from typing import Any import pandas as pd +from aqueduct.integrations.mongodb_integration import MongoDBIntegration from aqueduct import Client, Flow -from aqueduct.integrations.mongodb_integration import MongoDBIntegration from sdk.data_integration_tests.relational_data_validator import RelationalDataValidator from ..shared.validation import fetch_and_validate_saved_object_identifier diff --git a/integration_tests/sdk/data_integration_tests/mongo_db_test.py b/integration_tests/sdk/data_integration_tests/mongo_db_test.py index a96e19c4b..bc0730721 100644 --- a/integration_tests/sdk/data_integration_tests/mongo_db_test.py +++ b/integration_tests/sdk/data_integration_tests/mongo_db_test.py @@ -2,10 +2,10 @@ import pandas as pd import pytest - -from aqueduct import LoadUpdateMode, op from aqueduct.error import AqueductError from aqueduct.integrations.mongodb_integration import MongoDBIntegration + +from aqueduct import LoadUpdateMode, op from sdk.data_integration_tests.flow_manager import FlowManager from sdk.data_integration_tests.mongo_db_data_validator import MongoDBDataValidator from sdk.data_integration_tests.save import save diff --git a/integration_tests/sdk/data_integration_tests/relational_data_validator.py b/integration_tests/sdk/data_integration_tests/relational_data_validator.py index 0cfaa1903..8cdd7e582 100644 --- a/integration_tests/sdk/data_integration_tests/relational_data_validator.py +++ b/integration_tests/sdk/data_integration_tests/relational_data_validator.py @@ -2,12 +2,12 @@ from typing import Any, List, Tuple import pandas as pd - -from aqueduct import Client, Flow from aqueduct.constants.enums import LoadUpdateMode from aqueduct.integrations.sql_integration import RelationalDBIntegration from aqueduct.models.operators import RelationalDBLoadParams +from aqueduct import Client, Flow + from ..shared.validation import fetch_and_validate_saved_object_identifier diff --git a/integration_tests/sdk/data_integration_tests/relational_test.py b/integration_tests/sdk/data_integration_tests/relational_test.py index 9ee07dc6a..a7a8d6269 100644 --- a/integration_tests/sdk/data_integration_tests/relational_test.py +++ b/integration_tests/sdk/data_integration_tests/relational_test.py @@ -2,8 +2,6 @@ import pandas as pd import pytest - -from aqueduct import LoadUpdateMode, metric, op from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.artifacts.generic_artifact import GenericArtifact from aqueduct.constants.enums import ExecutionStatus @@ -11,6 +9,8 @@ from aqueduct.integrations.sql_integration import RelationalDBIntegration from aqueduct.models.operators import RelationalDBExtractParams +from aqueduct import LoadUpdateMode, metric, op + from ..shared.demo_db import demo_db_tables from ..shared.naming import generate_table_name from ..shared.relational import format_table_name diff --git a/integration_tests/sdk/data_integration_tests/s3_data_validator.py b/integration_tests/sdk/data_integration_tests/s3_data_validator.py index 320571d0e..93e0ca3d2 100644 --- a/integration_tests/sdk/data_integration_tests/s3_data_validator.py +++ b/integration_tests/sdk/data_integration_tests/s3_data_validator.py @@ -2,11 +2,11 @@ from typing import Any, Optional import pandas as pd +from aqueduct.constants.enums import ArtifactType +from aqueduct.integrations.s3_integration import S3Integration from PIL import Image from aqueduct import Client, Flow -from aqueduct.constants.enums import ArtifactType -from aqueduct.integrations.s3_integration import S3Integration from sdk.shared.validation import fetch_and_validate_saved_object_identifier diff --git a/integration_tests/sdk/data_integration_tests/s3_test.py b/integration_tests/sdk/data_integration_tests/s3_test.py index 36b49036d..fb20cdb9c 100644 --- a/integration_tests/sdk/data_integration_tests/s3_test.py +++ b/integration_tests/sdk/data_integration_tests/s3_test.py @@ -2,14 +2,14 @@ import pandas as pd import pytest -from PIL import Image - -from aqueduct import op from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.constants.enums import ArtifactType from aqueduct.error import AqueductError, InvalidUserArgumentException from aqueduct.flow import Flow from aqueduct.integrations.s3_integration import S3Integration +from PIL import Image + +from aqueduct import op from sdk.data_integration_tests.flow_manager import FlowManager from sdk.data_integration_tests.s3_data_validator import S3DataValidator from sdk.data_integration_tests.save import save diff --git a/integration_tests/sdk/setup_integration.py b/integration_tests/sdk/setup_integration.py index a1009a803..474b28157 100644 --- a/integration_tests/sdk/setup_integration.py +++ b/integration_tests/sdk/setup_integration.py @@ -5,8 +5,6 @@ import pandas as pd import yaml - -from aqueduct import Client, get_apikey from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.artifacts.table_artifact import TableArtifact from aqueduct.constants.enums import ArtifactType, ServiceType @@ -16,6 +14,8 @@ from aqueduct.integrations.s3_integration import S3Integration from aqueduct.integrations.sql_integration import RelationalDBIntegration from aqueduct.models.integration import Integration + +from aqueduct import Client, get_apikey from sdk.aqueduct_tests.save import save from sdk.shared.demo_db import demo_db_tables from sdk.shared.flow_helpers import delete_flow, publish_flow_test diff --git a/integration_tests/sdk/shared/flow_helpers.py b/integration_tests/sdk/shared/flow_helpers.py index 4adcc9dcc..4251221ac 100644 --- a/integration_tests/sdk/shared/flow_helpers.py +++ b/integration_tests/sdk/shared/flow_helpers.py @@ -2,11 +2,12 @@ import uuid from typing import Any, Dict, List, Optional, Union -import aqueduct -from aqueduct import Flow from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.constants.enums import ExecutionStatus +import aqueduct +from aqueduct import Flow + def publish_flow_test( client: aqueduct.Client, diff --git a/integration_tests/sdk/shared/utils.py b/integration_tests/sdk/shared/utils.py index d949b23a7..fb3b45aa2 100644 --- a/integration_tests/sdk/shared/utils.py +++ b/integration_tests/sdk/shared/utils.py @@ -2,13 +2,14 @@ import uuid from typing import Any, Dict, List, Optional, Union -import aqueduct -from aqueduct import Flow from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.constants.enums import ArtifactType, ExecutionStatus from aqueduct.integrations.s3_integration import S3Integration from aqueduct.integrations.sql_integration import RelationalDBIntegration +import aqueduct +from aqueduct import Flow + from .data_objects import DataObject diff --git a/integration_tests/sdk/shared/validation.py b/integration_tests/sdk/shared/validation.py index 455307ef8..a03fb202d 100644 --- a/integration_tests/sdk/shared/validation.py +++ b/integration_tests/sdk/shared/validation.py @@ -1,7 +1,8 @@ import uuid -from aqueduct import Client, Flow from aqueduct.models.integration import Integration + +from aqueduct import Client, Flow from sdk.shared.globals import artifact_id_to_saved_identifier """These helpers are shared across both integration test suites.""" diff --git a/integration_tests/sdk/shared/validator.py b/integration_tests/sdk/shared/validator.py index 02c2cc75d..3298d5c36 100644 --- a/integration_tests/sdk/shared/validator.py +++ b/integration_tests/sdk/shared/validator.py @@ -2,13 +2,13 @@ from typing import Any, List, Tuple import pandas as pd - -from aqueduct import Client, Flow from aqueduct.constants.enums import LoadUpdateMode from aqueduct.integrations.sql_integration import RelationalDBIntegration from aqueduct.models.integration import Integration from aqueduct.models.operators import RelationalDBLoadParams +from aqueduct import Client, Flow + from .globals import artifact_id_to_saved_identifier from .utils import extract diff --git a/manual_qa_tests/initialize.py b/manual_qa_tests/initialize.py index ad31174fc..7c068a809 100644 --- a/manual_qa_tests/initialize.py +++ b/manual_qa_tests/initialize.py @@ -2,6 +2,7 @@ from multiprocessing import Process import deploy_example +from aqueduct.constants.enums import NotificationLevel from notification import connect_slack from wait_for_flows import wait_for_all_flows_to_complete from workflows import ( @@ -17,7 +18,6 @@ ) import aqueduct as aq -from aqueduct.constants.enums import NotificationLevel # when adding new deployments, keep the order of `fail`, `warning`, and `succeed` # such that the UI would approximately show these workflows in reverse order. diff --git a/manual_qa_tests/notification.py b/manual_qa_tests/notification.py index 02e1d19d9..4e3004e86 100644 --- a/manual_qa_tests/notification.py +++ b/manual_qa_tests/notification.py @@ -1,7 +1,8 @@ -import aqueduct as aq from aqueduct.constants.enums import NotificationLevel from aqueduct.integrations.connect_config import ServiceType, SlackConfig +import aqueduct as aq + def connect_slack( client: aq.Client, diff --git a/manual_qa_tests/workflows/succeed_march_madness_dag_layout_test.py b/manual_qa_tests/workflows/succeed_march_madness_dag_layout_test.py index a8af35b01..9e716e804 100644 --- a/manual_qa_tests/workflows/succeed_march_madness_dag_layout_test.py +++ b/manual_qa_tests/workflows/succeed_march_madness_dag_layout_test.py @@ -1,6 +1,7 @@ -import aqueduct as aq from aqueduct.decorator import op +import aqueduct as aq + NAME = "succeed_march_madness_dag_layout_test" DESCRIPTION = """ * This test is here so that we can check for DAG positioning issues and edge overlaps. diff --git a/sdk/aqueduct/artifacts/bool_artifact.py b/sdk/aqueduct/artifacts/bool_artifact.py index 97c4e9aab..6987f028f 100644 --- a/sdk/aqueduct/artifacts/bool_artifact.py +++ b/sdk/aqueduct/artifacts/bool_artifact.py @@ -5,7 +5,6 @@ from typing import Any, Dict, Optional, Union import numpy as np - from aqueduct.artifacts import preview as artifact_utils from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.constants.enums import OperatorType diff --git a/sdk/aqueduct/artifacts/generic_artifact.py b/sdk/aqueduct/artifacts/generic_artifact.py index 917afc0df..574fad443 100644 --- a/sdk/aqueduct/artifacts/generic_artifact.py +++ b/sdk/aqueduct/artifacts/generic_artifact.py @@ -4,7 +4,6 @@ import uuid from typing import Any, Dict, Optional -from aqueduct import globals from aqueduct.artifacts import numeric_artifact from aqueduct.artifacts import preview as artifact_utils from aqueduct.artifacts import system_metric @@ -14,6 +13,8 @@ from aqueduct.models.dag import DAG from aqueduct.utils.utils import format_header_for_print +from aqueduct import globals + class GenericArtifact(BaseArtifact, system_metric.SystemMetricMixin): """This class represents a generic artifact within the flow's DAG. diff --git a/sdk/aqueduct/artifacts/numeric_artifact.py b/sdk/aqueduct/artifacts/numeric_artifact.py index 504472ebb..f1b11c7d3 100644 --- a/sdk/aqueduct/artifacts/numeric_artifact.py +++ b/sdk/aqueduct/artifacts/numeric_artifact.py @@ -5,8 +5,6 @@ from typing import Any, Callable, Dict, List, Optional import numpy as np - -from aqueduct import globals from aqueduct.artifacts import bool_artifact from aqueduct.artifacts import preview as artifact_utils from aqueduct.artifacts._create import create_metric_or_check_artifact @@ -36,6 +34,8 @@ from aqueduct.utils.naming import default_artifact_name_from_op_name from aqueduct.utils.utils import format_header_for_print, generate_uuid +from aqueduct import globals + class NumericArtifact(BaseArtifact): """This class represents a computed number within the flow's DAG. diff --git a/sdk/aqueduct/artifacts/preview.py b/sdk/aqueduct/artifacts/preview.py index 805d1a10f..b3a4a3f7e 100644 --- a/sdk/aqueduct/artifacts/preview.py +++ b/sdk/aqueduct/artifacts/preview.py @@ -3,7 +3,6 @@ import uuid from typing import Any, Dict, List, Optional -from aqueduct import globals from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.artifacts.create import to_artifact_class from aqueduct.constants.enums import ArtifactType, K8sClusterStatusType @@ -16,6 +15,8 @@ from aqueduct.utils.type_inference import infer_artifact_type from aqueduct.utils.utils import generate_engine_config +from aqueduct import globals + def preview_artifact( dag: DAG, target_artifact_id: uuid.UUID, parameters: Optional[Dict[str, Any]] = None diff --git a/sdk/aqueduct/artifacts/system_metric.py b/sdk/aqueduct/artifacts/system_metric.py index 6948d9054..3d2cc3fa6 100644 --- a/sdk/aqueduct/artifacts/system_metric.py +++ b/sdk/aqueduct/artifacts/system_metric.py @@ -1,7 +1,6 @@ import uuid from typing import Any, Dict, Union -from aqueduct import globals from aqueduct.artifacts import bool_artifact, numeric_artifact from aqueduct.artifacts import preview as artifact_utils from aqueduct.artifacts._create import create_metric_or_check_artifact @@ -13,6 +12,8 @@ from aqueduct.utils.naming import default_artifact_name_from_op_name from aqueduct.utils.utils import generate_uuid +from aqueduct import globals + class SystemMetricMixin: """A mixin class for the system_metric function. This is used by GenericArtifacts and TableArtifacts.""" diff --git a/sdk/aqueduct/artifacts/table_artifact.py b/sdk/aqueduct/artifacts/table_artifact.py index 5abdb1802..5b7c22296 100644 --- a/sdk/aqueduct/artifacts/table_artifact.py +++ b/sdk/aqueduct/artifacts/table_artifact.py @@ -5,9 +5,6 @@ from typing import Any, Dict, List, Optional import pandas as pd -from ruamel import yaml - -from aqueduct import globals from aqueduct.artifacts import bool_artifact, numeric_artifact from aqueduct.artifacts import preview as artifact_utils from aqueduct.artifacts import system_metric @@ -31,6 +28,9 @@ ) from aqueduct.utils.function_packaging import serialize_function from aqueduct.utils.utils import format_header_for_print +from ruamel import yaml + +from aqueduct import globals class TableArtifact(BaseArtifact, system_metric.SystemMetricMixin): diff --git a/sdk/aqueduct/backend/api_client.py b/sdk/aqueduct/backend/api_client.py index 30ce2c657..57b9c4881 100644 --- a/sdk/aqueduct/backend/api_client.py +++ b/sdk/aqueduct/backend/api_client.py @@ -4,8 +4,6 @@ from typing import IO, Any, DefaultDict, Dict, List, Optional, Tuple, Union import requests -from pkg_resources import get_distribution, parse_version - from aqueduct.constants.enums import ExecutionStatus, K8sClusterActionType, RuntimeType, ServiceType from aqueduct.error import ( AqueductError, @@ -35,6 +33,7 @@ SavedObjectUpdate, ) from aqueduct.utils.serialization import deserialize +from pkg_resources import get_distribution, parse_version from ..integrations.connect_config import DynamicK8sConfig, IntegrationConfig from .response_helpers import ( diff --git a/sdk/aqueduct/backend/response_helpers.py b/sdk/aqueduct/backend/response_helpers.py index 3e13ef3e8..df57f9068 100644 --- a/sdk/aqueduct/backend/response_helpers.py +++ b/sdk/aqueduct/backend/response_helpers.py @@ -3,8 +3,6 @@ import multipart import requests -from requests_toolbelt.multipart import decoder - from aqueduct.constants.enums import ExecutionStatus from aqueduct.error import AqueductError, InternalAqueductError from aqueduct.models.dag import DAG @@ -12,6 +10,7 @@ from aqueduct.models.operators import Operator from aqueduct.models.response_models import ArtifactResult, PreviewResponse from aqueduct.utils.utils import is_string_valid_uuid, print_logs +from requests_toolbelt.multipart import decoder def _parse_artifact_result_response(response: requests.Response) -> Dict[str, Any]: diff --git a/sdk/aqueduct/client.py b/sdk/aqueduct/client.py index 519862ccb..3d673c2e5 100644 --- a/sdk/aqueduct/client.py +++ b/sdk/aqueduct/client.py @@ -8,8 +8,6 @@ import __main__ as main import yaml - -from aqueduct import globals from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.artifacts.bool_artifact import BoolArtifact from aqueduct.artifacts.create import create_param_artifact @@ -66,6 +64,8 @@ generate_ui_url, ) +from aqueduct import globals + def global_config(config_dict: Dict[str, Any]) -> None: """Sets any global configuration variables in the current Aqueduct context. diff --git a/sdk/aqueduct/decorator.py b/sdk/aqueduct/decorator.py index 715d083ec..a670904f6 100644 --- a/sdk/aqueduct/decorator.py +++ b/sdk/aqueduct/decorator.py @@ -5,8 +5,6 @@ from typing import Any, Callable, Dict, List, Mapping, Optional, Union, cast import numpy as np - -from aqueduct import globals from aqueduct.artifacts._create import create_metric_or_check_artifact from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.artifacts.bool_artifact import BoolArtifact @@ -45,6 +43,8 @@ from aqueduct.utils.naming import default_artifact_name_from_op_name, sanitize_artifact_name from aqueduct.utils.utils import generate_engine_config, generate_uuid +from aqueduct import globals + OutputArtifactFunction = Callable[..., BaseArtifact] # For functions that can handle multiple outputs (eg. `op()`) diff --git a/sdk/aqueduct/error.py b/sdk/aqueduct/error.py index 471afd775..4494a977e 100644 --- a/sdk/aqueduct/error.py +++ b/sdk/aqueduct/error.py @@ -136,10 +136,12 @@ class ClientValidationError(Error): class RequirementsMissingError(Error): pass + # Exception raised when the user attempts to use an integration that has failed to connected. class IntegrationFailedToConnect(Error): pass + # Exception raised when the user attempts to use an integration that is still connecting. class IntegrationConnectionInProgress(Error): - pass \ No newline at end of file + pass diff --git a/sdk/aqueduct/flow.py b/sdk/aqueduct/flow.py index b5156a911..538feca6a 100644 --- a/sdk/aqueduct/flow.py +++ b/sdk/aqueduct/flow.py @@ -4,7 +4,6 @@ from collections import defaultdict from typing import DefaultDict, Dict, List, Optional, Union -from aqueduct import globals from aqueduct.error import InvalidUserArgumentException from aqueduct.flow_run import FlowRun from aqueduct.models.response_models import ( @@ -14,6 +13,8 @@ ) from aqueduct.utils.utils import format_header_for_print, generate_ui_url, parse_user_supplied_id +from aqueduct import globals + class Flow: """This class is a read-only handle to flow already in the system. diff --git a/sdk/aqueduct/flow_run.py b/sdk/aqueduct/flow_run.py index 7de33594c..7b173dcd1 100644 --- a/sdk/aqueduct/flow_run.py +++ b/sdk/aqueduct/flow_run.py @@ -3,7 +3,6 @@ from datetime import datetime from typing import Optional -from aqueduct import globals from aqueduct.artifacts import ( base_artifact, bool_artifact, @@ -22,6 +21,8 @@ print_logs, ) +from aqueduct import globals + class FlowRun: """This class is a read-only handle corresponding to a single workflow run in the system.""" diff --git a/sdk/aqueduct/github.py b/sdk/aqueduct/github.py index 888c3b3e9..d98292f82 100644 --- a/sdk/aqueduct/github.py +++ b/sdk/aqueduct/github.py @@ -1,6 +1,5 @@ from typing import List, Optional, Tuple -from aqueduct import globals from aqueduct.artifacts.table_artifact import TableArtifact from aqueduct.constants.enums import ( ArtifactType, @@ -19,6 +18,8 @@ ) from aqueduct.utils.function_packaging import DEFAULT_OP_METHOD_NAME, MODEL_FILE_NAME +from aqueduct import globals + def _get_operator_name(prefix: str, repo_url: str, branch: str, path: str) -> str: return f"{prefix}_{repo_url}_{branch}_{path}".replace("/", "_") diff --git a/sdk/aqueduct/integrations/connect_config.py b/sdk/aqueduct/integrations/connect_config.py index d7c727797..329fdf342 100644 --- a/sdk/aqueduct/integrations/connect_config.py +++ b/sdk/aqueduct/integrations/connect_config.py @@ -2,10 +2,9 @@ from enum import Enum from typing import Any, Dict, List, Optional, Union, cast -from pydantic import BaseModel, Extra, Field - from aqueduct.constants.enums import MetaEnum, NotificationLevel, ServiceType from aqueduct.error import InternalAqueductError, InvalidUserArgumentException +from pydantic import BaseModel, Extra, Field """Copied mostly over from `aqueduct_executor/operators/connectors/data/config.py` for now, please keep them in sync.""" diff --git a/sdk/aqueduct/integrations/dynamic_k8s_integration.py b/sdk/aqueduct/integrations/dynamic_k8s_integration.py index 9c5119895..5536d1c90 100644 --- a/sdk/aqueduct/integrations/dynamic_k8s_integration.py +++ b/sdk/aqueduct/integrations/dynamic_k8s_integration.py @@ -1,14 +1,14 @@ from typing import Dict, Union -from aqueduct.integrations.validation import validate_is_connected -from aqueduct import globals - from aqueduct.constants.enums import K8sClusterActionType, K8sClusterStatusType from aqueduct.error import InvalidIntegrationException, InvalidUserArgumentException from aqueduct.integrations.connect_config import DynamicK8sConfig +from aqueduct.integrations.validation import validate_is_connected from aqueduct.models.integration import Integration, IntegrationInfo from aqueduct.models.response_models import DynamicEngineStatusResponse +from aqueduct import globals + def parse_dynamic_k8s_config( config_delta: Union[Dict[str, Union[int, str]], DynamicK8sConfig] diff --git a/sdk/aqueduct/integrations/google_sheets_integration.py b/sdk/aqueduct/integrations/google_sheets_integration.py index acaef59f1..473090539 100644 --- a/sdk/aqueduct/integrations/google_sheets_integration.py +++ b/sdk/aqueduct/integrations/google_sheets_integration.py @@ -1,10 +1,9 @@ from typing import Optional -from aqueduct.integrations.validation import validate_is_connected - from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.artifacts.table_artifact import TableArtifact from aqueduct.constants.enums import ArtifactType, GoogleSheetsSaveMode +from aqueduct.integrations.validation import validate_is_connected from aqueduct.models.artifact import ArtifactMetadata from aqueduct.models.dag import DAG from aqueduct.models.integration import Integration, IntegrationInfo diff --git a/sdk/aqueduct/integrations/mongodb_integration.py b/sdk/aqueduct/integrations/mongodb_integration.py index f8504102c..9b74ac798 100644 --- a/sdk/aqueduct/integrations/mongodb_integration.py +++ b/sdk/aqueduct/integrations/mongodb_integration.py @@ -1,15 +1,13 @@ import json from typing import Any, Dict, List, Optional -from aqueduct.integrations.validation import validate_is_connected - -from aqueduct import globals from aqueduct.artifacts import preview as artifact_utils from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.artifacts.table_artifact import TableArtifact from aqueduct.constants.enums import ArtifactType, ExecutionMode, LoadUpdateMode from aqueduct.integrations.parameters import _validate_parameters from aqueduct.integrations.save import _save_artifact +from aqueduct.integrations.validation import validate_is_connected from aqueduct.models.artifact import ArtifactMetadata from aqueduct.models.dag import DAG from aqueduct.models.integration import Integration, IntegrationInfo @@ -24,6 +22,8 @@ from aqueduct.utils.naming import default_artifact_name_from_op_name, sanitize_artifact_name from aqueduct.utils.utils import generate_uuid +from aqueduct import globals + class MongoDBCollectionIntegration(Integration): _collection_name: str diff --git a/sdk/aqueduct/integrations/s3_integration.py b/sdk/aqueduct/integrations/s3_integration.py index a7fcc573a..afb43bb4d 100644 --- a/sdk/aqueduct/integrations/s3_integration.py +++ b/sdk/aqueduct/integrations/s3_integration.py @@ -1,12 +1,10 @@ import json from typing import List, Optional, Union -from aqueduct.integrations.validation import validate_is_connected - -from aqueduct import globals from aqueduct.artifacts import preview as artifact_utils from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.constants.enums import ArtifactType, ExecutionMode, S3TableFormat +from aqueduct.integrations.validation import validate_is_connected from aqueduct.models.artifact import ArtifactMetadata from aqueduct.models.dag import DAG from aqueduct.models.integration import Integration, IntegrationInfo @@ -20,6 +18,8 @@ from aqueduct.utils.dag_deltas import AddOperatorDelta, apply_deltas_to_dag from aqueduct.utils.utils import generate_uuid +from aqueduct import globals + from ..artifacts.create import to_artifact_class from ..error import InvalidUserArgumentException from ..utils.naming import default_artifact_name_from_op_name, sanitize_artifact_name diff --git a/sdk/aqueduct/integrations/salesforce_integration.py b/sdk/aqueduct/integrations/salesforce_integration.py index b00edffda..64d50a7ce 100644 --- a/sdk/aqueduct/integrations/salesforce_integration.py +++ b/sdk/aqueduct/integrations/salesforce_integration.py @@ -1,11 +1,10 @@ import uuid from typing import Optional -from aqueduct.integrations.validation import validate_is_connected - from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.artifacts.table_artifact import TableArtifact from aqueduct.constants.enums import ArtifactType, SalesforceExtractType +from aqueduct.integrations.validation import validate_is_connected from aqueduct.models.artifact import ArtifactMetadata from aqueduct.models.dag import DAG from aqueduct.models.integration import Integration, IntegrationInfo diff --git a/sdk/aqueduct/integrations/sql_integration.py b/sdk/aqueduct/integrations/sql_integration.py index 06a3e7f0c..e01e325e9 100644 --- a/sdk/aqueduct/integrations/sql_integration.py +++ b/sdk/aqueduct/integrations/sql_integration.py @@ -1,8 +1,6 @@ from typing import List, Optional, Union import pandas as pd -from aqueduct.integrations.validation import validate_is_connected - from aqueduct.artifacts.base_artifact import BaseArtifact from aqueduct.artifacts.preview import preview_artifact from aqueduct.artifacts.table_artifact import TableArtifact @@ -10,6 +8,7 @@ from aqueduct.error import InvalidUserActionException, InvalidUserArgumentException from aqueduct.integrations.parameters import _validate_builtin_expansions, _validate_parameters from aqueduct.integrations.save import _save_artifact +from aqueduct.integrations.validation import validate_is_connected from aqueduct.models.artifact import ArtifactMetadata from aqueduct.models.dag import DAG from aqueduct.models.integration import Integration, IntegrationInfo @@ -24,6 +23,8 @@ from aqueduct.utils.naming import default_artifact_name_from_op_name, sanitize_artifact_name from aqueduct.utils.utils import generate_uuid +from aqueduct import globals + LIST_TABLES_QUERY_PG = "SELECT tablename, tableowner FROM pg_catalog.pg_tables WHERE schemaname != 'pg_catalog' AND schemaname != 'information_schema';" LIST_TABLES_QUERY_SNOWFLAKE = "SELECT table_name AS \"tablename\", table_owner AS \"tableowner\" FROM information_schema.tables WHERE table_schema != 'INFORMATION_SCHEMA' AND table_type = 'BASE TABLE';" LIST_TABLES_QUERY_MYSQL = "SELECT table_name AS tablename FROM INFORMATION_SCHEMA.TABLES WHERE table_type = 'BASE TABLE' AND table_schema NOT IN ('mysql', 'sys', 'performance_schema');" @@ -281,4 +282,3 @@ def describe(self) -> None: print("(only first 5 tables are shown)") except: pass - diff --git a/sdk/aqueduct/integrations/validation.py b/sdk/aqueduct/integrations/validation.py index b44fd12ed..9d805853f 100644 --- a/sdk/aqueduct/integrations/validation.py +++ b/sdk/aqueduct/integrations/validation.py @@ -1,13 +1,19 @@ +from typing import Any, Callable + from aqueduct.utils.integration_validation import validate_integration_is_connected +AnyFunc = Callable[..., Any] + -def validate_is_connected(): +def validate_is_connected() -> Callable[[AnyFunc], AnyFunc]: """This decorator, when used on an Integration class method, ensures that the integration is connected before allowing the method to be called.""" - def decorator(method): - def wrapper(self, *args, **kwargs): + def decorator(method: AnyFunc) -> AnyFunc: + def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: validate_integration_is_connected(self.name(), self._metadata.exec_state) return method(*args, **kwargs) + return wrapper + return decorator diff --git a/sdk/aqueduct/models/artifact.py b/sdk/aqueduct/models/artifact.py index 79d0d67f0..bf0dbfc81 100644 --- a/sdk/aqueduct/models/artifact.py +++ b/sdk/aqueduct/models/artifact.py @@ -1,8 +1,7 @@ import uuid -from pydantic import BaseModel - from aqueduct.constants.enums import ArtifactType +from pydantic import BaseModel class ArtifactMetadata(BaseModel): diff --git a/sdk/aqueduct/models/config.py b/sdk/aqueduct/models/config.py index 6bfe06a51..0ddd469b9 100644 --- a/sdk/aqueduct/models/config.py +++ b/sdk/aqueduct/models/config.py @@ -1,14 +1,13 @@ import uuid from typing import Optional, Union -from pydantic import BaseModel - from aqueduct.constants.enums import RuntimeType from aqueduct.integrations.airflow_integration import AirflowIntegration from aqueduct.integrations.databricks_integration import DatabricksIntegration from aqueduct.integrations.k8s_integration import K8sIntegration from aqueduct.integrations.lambda_integration import LambdaIntegration from aqueduct.integrations.spark_integration import SparkIntegration +from pydantic import BaseModel class AqueductEngineConfig(BaseModel): diff --git a/sdk/aqueduct/models/dag.py b/sdk/aqueduct/models/dag.py index 4db1d8ea5..6cb25b49e 100644 --- a/sdk/aqueduct/models/dag.py +++ b/sdk/aqueduct/models/dag.py @@ -2,8 +2,6 @@ import uuid from typing import Any, Callable, Dict, List, Optional, Set -from pydantic import BaseModel - from aqueduct.constants.enums import ( ArtifactType, OperatorType, @@ -16,6 +14,7 @@ InternalAqueductError, InvalidUserActionException, ) +from pydantic import BaseModel from ..logger import logger from ..utils.naming import bump_artifact_suffix diff --git a/sdk/aqueduct/models/integration.py b/sdk/aqueduct/models/integration.py index e81dc4bbb..e91f86bb6 100644 --- a/sdk/aqueduct/models/integration.py +++ b/sdk/aqueduct/models/integration.py @@ -3,14 +3,11 @@ from abc import ABC from typing import Any, Dict, Optional -from aqueduct.constants.enums import ExecutionStatus -from aqueduct.error import IntegrationFailedToConnect, IntegrationConnectionInProgress +from aqueduct.constants.enums import ExecutionStatus, ServiceType from aqueduct.models.execution_state import ExecutionState from aqueduct.models.utils import human_readable_timestamp from pydantic import BaseModel -from aqueduct.constants.enums import ServiceType - class IntegrationInfo(BaseModel): id: uuid.UUID @@ -79,28 +76,6 @@ def name(self) -> str: def type(self) -> ServiceType: return self._metadata.service - def _validate_is_connected(self) -> None: - """Method used to determine if this integration was successfully connected to or not. - If not successfully connected (or pending), we will raise an Exception. - """ - # TODO(ENG-2813): Remove the assumption that a missing `exec_state` means success. - if self._metadata.exec_state is None or self._metadata.exec_state.status == ExecutionStatus.SUCCEEDED: - return - - if self._metadata.exec_state.status == ExecutionStatus.FAILED: - raise IntegrationFailedToConnect( - "The attempt to connect to %s has failed with error: %s\n%s\n\n " - "Please see the /integrations page on the UI for more details." % ( - self.name(), self._metadata.exec_state.error.tip, self._metadata.exec_state.error.context, - ) - ) - else: - raise IntegrationConnectionInProgress( - "We are still in the process of connecting to integration %s." - "Please see the /integrations page oen the UI for more details." - % self.name() - ) - def __hash__(self) -> int: """An integration is uniquely identified by its name. Ref: https://docs.python.org/3.5/reference/datamodel.html#object.__hash__ diff --git a/sdk/aqueduct/models/operators.py b/sdk/aqueduct/models/operators.py index 75fd1509f..c1ca0334e 100644 --- a/sdk/aqueduct/models/operators.py +++ b/sdk/aqueduct/models/operators.py @@ -1,8 +1,6 @@ import uuid from typing import List, Optional, Union -from pydantic import BaseModel, Extra - from aqueduct.constants.enums import ( ArtifactType, CheckSeverity, @@ -20,6 +18,7 @@ from aqueduct.error import AqueductError, UnsupportedFeatureException from aqueduct.models.config import EngineConfig from aqueduct.models.integration import IntegrationInfo +from pydantic import BaseModel, Extra class GithubMetadata(BaseModel): diff --git a/sdk/aqueduct/models/response_models.py b/sdk/aqueduct/models/response_models.py index 31acee8b5..c22d8e60b 100644 --- a/sdk/aqueduct/models/response_models.py +++ b/sdk/aqueduct/models/response_models.py @@ -2,8 +2,6 @@ from datetime import datetime from typing import Dict, List, Optional -from pydantic import BaseModel - from aqueduct.constants.enums import ( ArtifactType, ExecutionStatus, @@ -15,6 +13,7 @@ from aqueduct.models.execution_state import ExecutionState from aqueduct.models.operators import LoadSpec, Operator, OperatorSpec from aqueduct.models.utils import human_readable_timestamp +from pydantic import BaseModel class ArtifactResult(BaseModel): diff --git a/sdk/aqueduct/tests/decorator_test.py b/sdk/aqueduct/tests/decorator_test.py index 6a95d650d..546cbbaa4 100644 --- a/sdk/aqueduct/tests/decorator_test.py +++ b/sdk/aqueduct/tests/decorator_test.py @@ -3,7 +3,6 @@ import zipfile import pytest - from aqueduct.error import ( InvalidDependencyFilePath, RequirementsMissingError, diff --git a/sdk/aqueduct/tests/decorators_with_without_parentheses_test.py b/sdk/aqueduct/tests/decorators_with_without_parentheses_test.py index 10e6442aa..c2295c055 100644 --- a/sdk/aqueduct/tests/decorators_with_without_parentheses_test.py +++ b/sdk/aqueduct/tests/decorators_with_without_parentheses_test.py @@ -1,8 +1,6 @@ from unittest.mock import MagicMock import pandas as pd - -from aqueduct import globals from aqueduct.artifacts.bool_artifact import BoolArtifact from aqueduct.artifacts.numeric_artifact import NumericArtifact from aqueduct.artifacts.table_artifact import TableArtifact @@ -11,6 +9,8 @@ from aqueduct.tests.utils import construct_mocked_preview, default_table_artifact from aqueduct.utils.function_packaging import delete_zip_folder_and_file +from aqueduct import globals + def test_decorators_with_without_parentheses(): inp = default_table_artifact() diff --git a/sdk/aqueduct/tests/flow_test.py b/sdk/aqueduct/tests/flow_test.py index 04947f4a4..f0d3e5cf7 100644 --- a/sdk/aqueduct/tests/flow_test.py +++ b/sdk/aqueduct/tests/flow_test.py @@ -2,7 +2,6 @@ from typing import Tuple import pytest - from aqueduct.error import InvalidUserArgumentException from aqueduct.utils.utils import find_flow_with_user_supplied_id_and_name diff --git a/sdk/aqueduct/tests/metric_test.py b/sdk/aqueduct/tests/metric_test.py index 3c04cc2f2..4ab0fd575 100644 --- a/sdk/aqueduct/tests/metric_test.py +++ b/sdk/aqueduct/tests/metric_test.py @@ -3,7 +3,6 @@ from io import StringIO from unittest.mock import MagicMock -from aqueduct import globals from aqueduct.artifacts.numeric_artifact import NumericArtifact from aqueduct.constants.enums import ArtifactType, SerializationType from aqueduct.decorator import check, metric @@ -11,6 +10,8 @@ from aqueduct.utils.function_packaging import delete_zip_folder_and_file from aqueduct.utils.utils import generate_uuid +from aqueduct import globals + metric_op_name = "metric_fn" description = f"{metric_op_name} description" metric_artifact_name = f"{metric_op_name} artifact" diff --git a/sdk/aqueduct/tests/serialization_test.py b/sdk/aqueduct/tests/serialization_test.py index f1b000fc0..f0c2b3a02 100644 --- a/sdk/aqueduct/tests/serialization_test.py +++ b/sdk/aqueduct/tests/serialization_test.py @@ -2,8 +2,6 @@ import uuid import cloudpickle as pickle -from PIL import Image - from aqueduct.constants.enums import ( ArtifactType, ExecutionStatus, @@ -47,6 +45,7 @@ serialize_val, ) from aqueduct.utils.utils import generate_uuid +from PIL import Image def test_artifact_serialization(): diff --git a/sdk/aqueduct/tests/utils.py b/sdk/aqueduct/tests/utils.py index e957cf21b..334557ab1 100644 --- a/sdk/aqueduct/tests/utils.py +++ b/sdk/aqueduct/tests/utils.py @@ -3,8 +3,6 @@ from typing import List, Optional import pandas as pd - -from aqueduct import globals from aqueduct.artifacts.table_artifact import TableArtifact from aqueduct.constants.enums import ( ArtifactType, @@ -33,6 +31,8 @@ from aqueduct.models.response_models import ArtifactResult, PreviewResponse from aqueduct.utils.utils import generate_uuid +from aqueduct import globals + def generate_uuids(num: int) -> List[uuid.UUID]: return [generate_uuid() for _ in range(num)] diff --git a/sdk/aqueduct/utils/function_packaging.py b/sdk/aqueduct/utils/function_packaging.py index e1ef840ef..cc9eb514d 100644 --- a/sdk/aqueduct/utils/function_packaging.py +++ b/sdk/aqueduct/utils/function_packaging.py @@ -9,7 +9,6 @@ import cloudpickle as pickle import pkg_resources - from aqueduct.error import ( InternalAqueductError, InvalidDependencyFilePath, diff --git a/sdk/aqueduct/utils/integration_validation.py b/sdk/aqueduct/utils/integration_validation.py index 1d4bc6094..d6d0373ac 100644 --- a/sdk/aqueduct/utils/integration_validation.py +++ b/sdk/aqueduct/utils/integration_validation.py @@ -1,8 +1,7 @@ from typing import Optional -from aqueduct.error import IntegrationFailedToConnect, IntegrationConnectionInProgress - from aqueduct.constants.enums import ExecutionStatus +from aqueduct.error import IntegrationConnectionInProgress, IntegrationFailedToConnect from aqueduct.models.execution_state import ExecutionState @@ -15,16 +14,18 @@ def validate_integration_is_connected(name: str, exec_state: Optional[ExecutionS return if exec_state.status == ExecutionStatus.FAILED: + assert exec_state.error is not None raise IntegrationFailedToConnect( "Cannot use integration %s because it has not been successfully connected to: " - "%s\n%s\n\n Please see the /integrations page on the UI for more details." % ( - name, exec_state.error.tip, exec_state.error.context, + "%s\n%s\n\n Please see the /integrations page on the UI for more details." + % ( + name, + exec_state.error.tip, + exec_state.error.context, ) ) else: raise IntegrationConnectionInProgress( "Cannot use integration %s because it is still in the process of connecting." - "Please see the /integrations page on the UI for more details." - % name + "Please see the /integrations page on the UI for more details." % name ) - \ No newline at end of file diff --git a/sdk/aqueduct/utils/serialization.py b/sdk/aqueduct/utils/serialization.py index f7522644e..e07cb7a4b 100644 --- a/sdk/aqueduct/utils/serialization.py +++ b/sdk/aqueduct/utils/serialization.py @@ -6,10 +6,6 @@ import cloudpickle as pickle import pandas as pd -from bson import json_util as bson_json_util -from PIL import Image -from pydantic import BaseModel - from aqueduct.constants.enums import ( ArtifactType, LocalDataSerializationType, @@ -19,6 +15,9 @@ ) from aqueduct.utils.local_data import _convert_to_local_data_table_format from aqueduct.utils.type_inference import infer_artifact_type +from bson import json_util as bson_json_util +from PIL import Image +from pydantic import BaseModel from .format import DEFAULT_ENCODING from .function_packaging import _make_temp_dir diff --git a/sdk/aqueduct/utils/type_inference.py b/sdk/aqueduct/utils/type_inference.py index 18eb048f1..dc1f8b0c0 100644 --- a/sdk/aqueduct/utils/type_inference.py +++ b/sdk/aqueduct/utils/type_inference.py @@ -4,11 +4,10 @@ import cloudpickle as pickle import numpy as np +from aqueduct.constants.enums import ArtifactType from pandas import DataFrame from PIL import Image -from aqueduct.constants.enums import ArtifactType - from .format import DEFAULT_ENCODING diff --git a/sdk/aqueduct/utils/utils.py b/sdk/aqueduct/utils/utils.py index 8f1820001..7bd0f71b4 100644 --- a/sdk/aqueduct/utils/utils.py +++ b/sdk/aqueduct/utils/utils.py @@ -1,10 +1,6 @@ import uuid from typing import Any, Dict, List, Optional, Tuple, Union -from aqueduct.utils.integration_validation import validate_integration_is_connected - -from croniter import croniter - from aqueduct.constants.enums import ArtifactType, RuntimeType, ServiceType, TriggerType from aqueduct.error import * from aqueduct.models.config import ( @@ -18,6 +14,8 @@ from aqueduct.models.dag import Schedule from aqueduct.models.integration import IntegrationInfo from aqueduct.models.operators import ParamSpec +from aqueduct.utils.integration_validation import validate_integration_is_connected +from croniter import croniter from ..models.execution_state import Logs from .serialization import artifact_type_to_serialization_type, serialize_val diff --git a/src/python/aqueduct_executor/migrators/artifact_migration_000016/execute.py b/src/python/aqueduct_executor/migrators/artifact_migration_000016/execute.py index 417223a22..e568b0152 100644 --- a/src/python/aqueduct_executor/migrators/artifact_migration_000016/execute.py +++ b/src/python/aqueduct_executor/migrators/artifact_migration_000016/execute.py @@ -2,13 +2,12 @@ from typing import Any import cloudpickle as pickle +from aqueduct.utils.type_inference import infer_artifact_type from aqueduct_executor.migrators.artifact_migration_000016.spec import MigrationSpec from aqueduct_executor.operators.utils.enums import ArtifactType, SerializationType from aqueduct_executor.operators.utils.storage.parse import parse_storage from aqueduct_executor.operators.utils.storage.storage import Storage -from aqueduct.utils.type_inference import infer_artifact_type - # The variable definition and type mapping logic below are partially # duplicated from aqueduct_executor.operators.utils.utils diff --git a/src/python/aqueduct_executor/migrators/parameter_val_type_inference_000019/execute.py b/src/python/aqueduct_executor/migrators/parameter_val_type_inference_000019/execute.py index 952de75d3..5adbb51ec 100644 --- a/src/python/aqueduct_executor/migrators/parameter_val_type_inference_000019/execute.py +++ b/src/python/aqueduct_executor/migrators/parameter_val_type_inference_000019/execute.py @@ -1,11 +1,10 @@ import base64 import json +from aqueduct.utils.type_inference import infer_artifact_type from aqueduct_executor.migrators.parameter_val_type_inference_000019 import serialize from aqueduct_executor.operators.utils.enums import SerializationType -from aqueduct.utils.type_inference import infer_artifact_type - def run_type_inference_and_encode(json_val: str) -> None: """ diff --git a/src/python/aqueduct_executor/operators/connectors/data/extract.py b/src/python/aqueduct_executor/operators/connectors/data/extract.py index 66470008a..7d78d76b4 100644 --- a/src/python/aqueduct_executor/operators/connectors/data/extract.py +++ b/src/python/aqueduct_executor/operators/connectors/data/extract.py @@ -3,12 +3,11 @@ import uuid from typing import Any, Dict, List, Optional, Union +from aqueduct.integrations.parameters import BUILT_IN_EXPANSIONS, TAG_PATTERN from aqueduct_executor.operators.connectors.data import common, models from aqueduct_executor.operators.utils.enums import ArtifactType from pydantic import parse_obj_as -from aqueduct.integrations.parameters import BUILT_IN_EXPANSIONS, TAG_PATTERN - # The TAG for 'previous table' when the user specifies a chained query. PREV_TABLE_TAG = "$" diff --git a/src/python/aqueduct_executor/operators/connectors/data/s3.py b/src/python/aqueduct_executor/operators/connectors/data/s3.py index 224e1743c..bfa1f887d 100644 --- a/src/python/aqueduct_executor/operators/connectors/data/s3.py +++ b/src/python/aqueduct_executor/operators/connectors/data/s3.py @@ -3,6 +3,7 @@ import numpy as np import pandas as pd +from aqueduct.utils.serialization import deserialize from aqueduct_executor.operators.connectors.data import connector, extract, load from aqueduct_executor.operators.connectors.data.config import S3Config from aqueduct_executor.operators.connectors.data.s3_serialization import ( @@ -18,8 +19,6 @@ from aqueduct_executor.operators.utils.utils import delete_object from botocore.client import ClientError -from aqueduct.utils.serialization import deserialize - class S3Connector(connector.DataConnector): def __init__(self, config: S3Config): diff --git a/src/python/aqueduct_executor/operators/connectors/data/s3_serialization.py b/src/python/aqueduct_executor/operators/connectors/data/s3_serialization.py index 0f03607a2..ea0ef599d 100644 --- a/src/python/aqueduct_executor/operators/connectors/data/s3_serialization.py +++ b/src/python/aqueduct_executor/operators/connectors/data/s3_serialization.py @@ -4,13 +4,12 @@ import cloudpickle as pickle import pandas as pd -from aqueduct_executor.operators.connectors.data.common import S3TableFormat -from aqueduct_executor.operators.utils.enums import ArtifactType -from PIL import Image - from aqueduct.constants.enums import S3SerializationType from aqueduct.utils.serialization import PickleableCollectionSerializationFormat from aqueduct.utils.type_inference import infer_artifact_type +from aqueduct_executor.operators.connectors.data.common import S3TableFormat +from aqueduct_executor.operators.utils.enums import ArtifactType +from PIL import Image _DEFAULT_JSON_ENCODING = "utf8" _DEFAULT_IMAGE_FORMAT = "jpeg" diff --git a/src/python/aqueduct_executor/operators/function_executor/execute.py b/src/python/aqueduct_executor/operators/function_executor/execute.py index abbf2e8cd..f4852b818 100644 --- a/src/python/aqueduct_executor/operators/function_executor/execute.py +++ b/src/python/aqueduct_executor/operators/function_executor/execute.py @@ -9,6 +9,8 @@ import numpy as np import pandas as pd +from aqueduct.utils.serialization import check_and_fetch_pickled_collection_format +from aqueduct.utils.type_inference import infer_artifact_type from aqueduct_executor.operators.function_executor import extract_function, get_extract_path from aqueduct_executor.operators.function_executor.spec import FunctionSpec from aqueduct_executor.operators.function_executor.utils import OP_DIR @@ -37,9 +39,6 @@ from aqueduct_executor.operators.utils.timer import Timer from aqueduct_executor.operators.utils.utils import time_it -from aqueduct.utils.serialization import check_and_fetch_pickled_collection_format -from aqueduct.utils.type_inference import infer_artifact_type - def get_py_import_path(spec: FunctionSpec) -> str: """ diff --git a/src/python/aqueduct_executor/operators/param_executor/execute.py b/src/python/aqueduct_executor/operators/param_executor/execute.py index e25ab572e..5b5a4a113 100644 --- a/src/python/aqueduct_executor/operators/param_executor/execute.py +++ b/src/python/aqueduct_executor/operators/param_executor/execute.py @@ -1,5 +1,7 @@ import sys +from aqueduct.utils.serialization import deserialize +from aqueduct.utils.type_inference import infer_artifact_type from aqueduct_executor.operators.param_executor.spec import ParamSpec from aqueduct_executor.operators.utils import utils from aqueduct_executor.operators.utils.enums import ExecutionStatus, FailureType @@ -13,9 +15,6 @@ ) from aqueduct_executor.operators.utils.storage.parse import parse_storage -from aqueduct.utils.serialization import deserialize -from aqueduct.utils.type_inference import infer_artifact_type - def run(spec: ParamSpec) -> None: """Parameter operators are unique in that the output content is expected to have already been populated correctly. diff --git a/src/python/aqueduct_executor/operators/spark/execute_function.py b/src/python/aqueduct_executor/operators/spark/execute_function.py index ede6781a4..9631c0005 100644 --- a/src/python/aqueduct_executor/operators/spark/execute_function.py +++ b/src/python/aqueduct_executor/operators/spark/execute_function.py @@ -9,6 +9,7 @@ import numpy as np import pandas as pd +from aqueduct.utils.type_inference import infer_artifact_type from aqueduct_executor.operators.function_executor import extract_function, get_extract_path from aqueduct_executor.operators.function_executor.execute import ( cleanup, @@ -43,8 +44,6 @@ from aqueduct_executor.operators.utils.timer import Timer from pyspark.sql import SparkSession, dataframe -from aqueduct.utils.type_inference import infer_artifact_type - def _infer_artifact_type_spark(value: Any) -> Any: if isinstance(value, dataframe.DataFrame): diff --git a/src/python/aqueduct_executor/operators/spark/utils.py b/src/python/aqueduct_executor/operators/spark/utils.py index 3b6003f76..3014171c4 100644 --- a/src/python/aqueduct_executor/operators/spark/utils.py +++ b/src/python/aqueduct_executor/operators/spark/utils.py @@ -3,6 +3,12 @@ from typing import Any, Callable, Dict, List, Optional, Tuple import pandas as pd +from aqueduct.utils.format import DEFAULT_ENCODING +from aqueduct.utils.serialization import ( + artifact_type_to_serialization_type, + deserialize, + serialize_val, +) from aqueduct_executor.operators.utils.enums import ( ArtifactType, ExecutionStatus, @@ -30,13 +36,6 @@ ) from pyspark.sql import SparkSession -from aqueduct.utils.format import DEFAULT_ENCODING -from aqueduct.utils.serialization import ( - artifact_type_to_serialization_type, - deserialize, - serialize_val, -) - def read_artifacts_spark( storage: Storage, diff --git a/src/python/aqueduct_executor/operators/utils/utils.py b/src/python/aqueduct_executor/operators/utils/utils.py index 46d3e67d0..7dd3e3e92 100644 --- a/src/python/aqueduct_executor/operators/utils/utils.py +++ b/src/python/aqueduct_executor/operators/utils/utils.py @@ -4,6 +4,12 @@ from typing import Any, Callable, Dict, List, Optional, Tuple import pandas as pd +from aqueduct.utils.format import DEFAULT_ENCODING +from aqueduct.utils.serialization import ( + artifact_type_to_serialization_type, + deserialize, + serialize_val, +) from aqueduct_executor.operators.utils.enums import ( ArtifactType, ExecutionStatus, @@ -23,13 +29,6 @@ from aqueduct_executor.operators.utils.saved_object_delete import SavedObjectDelete from aqueduct_executor.operators.utils.storage.storage import Storage -from aqueduct.utils.format import DEFAULT_ENCODING -from aqueduct.utils.serialization import ( - artifact_type_to_serialization_type, - deserialize, - serialize_val, -) - _RUNTIME_SEC_METRIC_NAME = "runtime" _MAX_MEMORY_MB_METRIC_NAME = "max_memory" _METADATA_SCHEMA_KEY = "schema"