diff --git a/app/bulk_persistence/__init__.py b/app/bulk_persistence/__init__.py index 2d0a6fc95d5fbdbe9b1d9e1f2b0be56ede03a3e5..4bba26a8f4f9912e0664e886e44ea467245ccc3f 100644 --- a/app/bulk_persistence/__init__.py +++ b/app/bulk_persistence/__init__.py @@ -13,6 +13,8 @@ # limitations under the License. from .bulk_uri import BulkURI +from .dask.dask_bulk_storage import DaskBulkStorage +from .dask.dask_bulk_storage_local import make_local_dask_bulk_storage from .dataframe_persistence import create_and_store_dataframe, get_dataframe, download_bulk from .dataframe_serializer import DataframeSerializerAsync, DataframeSerializerSync from .json_orient import JSONOrient diff --git a/app/bulk_persistence/dask/dask_bulk_storage.py b/app/bulk_persistence/dask/dask_bulk_storage.py index 1fce29896fd88882859a1f3bdf29d8dbfe5c8cf9..17d04fe1c356c9f41b8d66a6c795c7b49e858793 100644 --- a/app/bulk_persistence/dask/dask_bulk_storage.py +++ b/app/bulk_persistence/dask/dask_bulk_storage.py @@ -37,7 +37,7 @@ from .traces import map_with_trace, submit_with_trace, trace_attributes_root_spa from .utils import (WDMS_INDEX_NAME, by_pairs, do_merge, join_dataframes, worker_capture_timing_handlers, get_num_rows, set_index, index_union) from ..dataframe_validators import is_reserved_column_name, DataFrameValidationFunc -from .. import DataframeSerializerSync +from ..dataframe_serializer import DataframeSerializerSync from . import storage_path_builder as pathBuilder from . import session_file_meta as session_meta from ..bulk_id import new_bulk_id diff --git a/app/consistency/trajectory_consistency.py b/app/consistency/trajectory_consistency.py index c54a76d6afacb4bf0dcd7c0a684a6f00af18ee31..982ee0487bf03f0af8c4f951160b54f5eedd1e5a 100644 --- a/app/consistency/trajectory_consistency.py +++ b/app/consistency/trajectory_consistency.py @@ -8,7 +8,8 @@ from odes_storage.models import Record from app.model.osdu_model import WellboreTrajectory110 from app.helper.traces import with_trace from app.bulk_persistence.consistency_checks import ConsistencyException, DataConsistencyChecks -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage, BulkRecordNotFound +from app.bulk_persistence.dask.dask_bulk_storage import BulkRecordNotFound +from app.bulk_persistence import DaskBulkStorage from app.bulk_persistence.dask.traces import submit_with_trace from app.model.model_utils import from_record from app.context import get_ctx diff --git a/app/consistency/welllog_consistency.py b/app/consistency/welllog_consistency.py index e33d55a07747b36c0835eca65e386facf87d8ecf..650d7c7a53cbe0bc67302dd7cdeb6d9a1ece9d0e 100644 --- a/app/consistency/welllog_consistency.py +++ b/app/consistency/welllog_consistency.py @@ -8,7 +8,8 @@ from odes_storage.models import Record from app.helper.traces import with_trace from app.bulk_persistence.consistency_checks import ConsistencyException, DataConsistencyChecks -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage, BulkRecordNotFound +from app.bulk_persistence.dask.dask_bulk_storage import BulkRecordNotFound +from app.bulk_persistence import DaskBulkStorage from app.bulk_persistence.dask.traces import submit_with_trace from app.model.model_utils import from_record from app.model.osdu_model import WellLog110 diff --git a/app/injector/aws_injector.py b/app/injector/aws_injector.py index a35f2711d03c20b3f3bff294cdefe793387734dc..0e78e05b01321153dc442619c2c79f0a5e0cd067 100644 --- a/app/injector/aws_injector.py +++ b/app/injector/aws_injector.py @@ -16,8 +16,7 @@ from osdu.core.api.storage.blob_storage_base import BlobStorageBase from osdu_aws.storage.storage_aws import AwsStorage from osdu_aws.storage.dask_storage_parameters import get_dask_storage_parameters as aws_parameters from app.context import Context -from app.bulk_persistence import resolve_tenant -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage +from app.bulk_persistence import resolve_tenant, DaskBulkStorage from .app_injector import AppInjector, AppInjectorModule from app.conf import Config diff --git a/app/injector/az_injector.py b/app/injector/az_injector.py index da8e797687a547658b4e11e64ef81078dc8e4720..ea12c919ecafca4acac48a9e478f3bb0da0ca5fe 100644 --- a/app/injector/az_injector.py +++ b/app/injector/az_injector.py @@ -13,7 +13,7 @@ # limitations under the License. from app.bulk_persistence import resolve_tenant -from app.bulk_persistence.dask.dask_bulk_storage import (DaskBulkStorage) +from app.bulk_persistence import DaskBulkStorage from app.context import Context from osdu.core.api.storage.blob_storage_base import BlobStorageBase from osdu_az.storage.blob_storage_az import AzureAioBlobStorage diff --git a/app/injector/gcp_injector.py b/app/injector/gcp_injector.py index 003b5eb0b39a5301f2b17fcf0c46056e7da8d8ff..d1fc5b54b1e4b9edc19e1eba6d26a4a0328f79ac 100644 --- a/app/injector/gcp_injector.py +++ b/app/injector/gcp_injector.py @@ -18,8 +18,7 @@ from osdu_gcp.storage.blob_storage_gcp import GCloudAioStorage from app.utils import get_http_client_session from .app_injector import AppInjector, AppInjectorModule from app.context import Context -from app.bulk_persistence import resolve_tenant -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage +from app.bulk_persistence import resolve_tenant, DaskBulkStorage from osdu_gcp.storage.dask_storage_parameters import get_dask_storage_parameters as gcp_parameters diff --git a/app/injector/ibm_injector.py b/app/injector/ibm_injector.py index 444b4a06031a9f0c21e6a20ac94bfec46f0b6f6a..b856a013d7d9d846918224eb3f10c98c8bd0b03d 100644 --- a/app/injector/ibm_injector.py +++ b/app/injector/ibm_injector.py @@ -4,8 +4,7 @@ from osdu_ibm.storage.blob_storage_ibm import IBMObjectStorage from app.utils import get_http_client_session from app.context import Context from .app_injector import AppInjector, AppInjectorModule -from app.bulk_persistence import resolve_tenant -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage +from app.bulk_persistence import resolve_tenant, DaskBulkStorage from osdu_ibm.storage.dask_storage_parameters import DaskStorageParametersFactoryIBM diff --git a/app/injector/main_injector.py b/app/injector/main_injector.py index 681752fc61ae4b60ab79486807920aad9c28e337..2125bcb51fd2fdb7accc537d4d88c421dcb34820 100644 --- a/app/injector/main_injector.py +++ b/app/injector/main_injector.py @@ -27,8 +27,7 @@ from app.clients.search_service_client import SearchServiceClient from app.clients import make_search_client, make_storage_record_client from app.bulk_persistence import SessionsStorage -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage -from app.bulk_persistence.dask.dask_bulk_storage_local import make_local_dask_bulk_storage +from app.bulk_persistence import DaskBulkStorage, make_local_dask_bulk_storage class MainInjector(AppInjectorModule): diff --git a/app/routers/bulk/bulk_routes.py b/app/routers/bulk/bulk_routes.py index e07f6c535e668338c5d31ad4d2cfadb5c0aec9df..d0b1da2770dd3fd9e14e5a3fcd76788ba0225ef6 100644 --- a/app/routers/bulk/bulk_routes.py +++ b/app/routers/bulk/bulk_routes.py @@ -62,7 +62,7 @@ from app.bulk_persistence.dataframe_validators import (auto_cast_columns_to_stri DataFrameValidationFunc, no_validation) from app.bulk_persistence import JSONOrient, get_dataframe, download_bulk -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage +from app.bulk_persistence import DaskBulkStorage from app.bulk_persistence.dask.errors import BulkError, BulkRecordNotFound, FilterError, TooManyColumnsRequested from app.bulk_persistence.mime_types import MimeTypes, MimeType from app.bulk_persistence.dask.traces import trace_dataframe_attributes, trace_attributes_root_span diff --git a/app/routers/bulk/utils.py b/app/routers/bulk/utils.py index 02b5b50c9dce864381e0858c2786af1b5f6b026e..8af7495d882213da3682d5b66bfa35bcb6ed98c7 100644 --- a/app/routers/bulk/utils.py +++ b/app/routers/bulk/utils.py @@ -13,7 +13,7 @@ from pyarrow.lib import ArrowInvalid from app.bulk_persistence.dask.errors import FilterError, internal_bulk_exceptions, BulkCurvesNotFound from app.bulk_persistence.dask.traces import trace_dataframe_attributes -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage +from app.bulk_persistence import DaskBulkStorage from app.bulk_persistence.dataframe_validators import auto_cast_columns_to_string, columns_type_must_be_string, \ no_validation, DataFrameValidationFunc from app.bulk_persistence import DataframeSerializerAsync diff --git a/tests/unit/routers/chunking_test.py b/tests/unit/routers/chunking_test.py index 58e72516d33862015d2cdf034745e140fd8fe52d..23d68907fa9589751555d0c5286fc53ee07206ba 100644 --- a/tests/unit/routers/chunking_test.py +++ b/tests/unit/routers/chunking_test.py @@ -9,13 +9,11 @@ import pyarrow as pa import pyarrow.parquet as pq import pytest from app.auth.auth import require_opendes_authorized_user -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage -from app.bulk_persistence.dask.dask_bulk_storage_local import make_local_dask_bulk_storage +from app.bulk_persistence import DaskBulkStorage, make_local_dask_bulk_storage from app.clients import StorageRecordServiceClient from app.clients.storage_service_blob_storage import StorageRecordServiceBlobStorage from app.helper import traces from app.middleware import require_data_partition_id -from app.bulk_persistence import SessionsStorage, SessionState from app.conf import Config from app.wdms_app import app_injector, wdms_app diff --git a/tests/unit/routers/ddms_v3/common_ddms_v3_test.py b/tests/unit/routers/ddms_v3/common_ddms_v3_test.py index 21512e7fd75e60d7242b48ffe3d0ffaa42c9b954..5c8aacc27c040ae34bbc72abd0a450168defc341 100644 --- a/tests/unit/routers/ddms_v3/common_ddms_v3_test.py +++ b/tests/unit/routers/ddms_v3/common_ddms_v3_test.py @@ -18,13 +18,11 @@ from odes_storage import UnexpectedResponse import pandas as pd import pytest from app.auth.auth import require_opendes_authorized_user -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage -from app.bulk_persistence.dask.dask_bulk_storage_local import make_local_dask_bulk_storage +from app.bulk_persistence import DaskBulkStorage, make_local_dask_bulk_storage, SessionsStorage from app.clients import SearchServiceClient, StorageRecordServiceClient from app.helper import traces from app.middleware import require_data_partition_id from app.model.osdu_model import Well, Wellbore -from app.bulk_persistence import SessionsStorage from app.wdms_app import app_injector, wdms_app from fastapi import status from fastapi.testclient import TestClient diff --git a/tests/unit/routers/ddms_v3/delete_ddms_v3_test.py b/tests/unit/routers/ddms_v3/delete_ddms_v3_test.py index 9ddd4605681ef3e15d57012314eda24ef1972c6c..c910bb73293ac5bb2c7c230c88f476f6bef1380a 100644 --- a/tests/unit/routers/ddms_v3/delete_ddms_v3_test.py +++ b/tests/unit/routers/ddms_v3/delete_ddms_v3_test.py @@ -22,7 +22,7 @@ from opencensus.trace import base_exporter from osdu.core.api.storage.blob_storage_base import BlobStorageBase -from app.bulk_persistence.dask.dask_bulk_storage import DaskBulkStorage +from app.bulk_persistence import DaskBulkStorage from app.clients import StorageRecordServiceClient from odes_storage.models import RecordVersions