tracdap-runtime 0.6.0rc1__py3-none-any.whl → 0.6.0rc3__py3-none-any.whl

tracdap/rt/_plugins/storage_aws.py

@@ -31,19 +31,19 @@ from . import _helpers
 
 try:
     # AWS SDK
-    import boto3  # noqa
-    import botocore.response  # noqa
-    import botocore.exceptions as aws_ex  # noqa
-    boto_available = True
+    import boto3
+    import botocore.response
+    import botocore.exceptions as aws_ex
+    __aws_available = True
 except ImportError:
-    boto_available = False
+    boto3 = None
+    botocore = None
+    aws_ex = None
+    __aws_available = False
 
 
 class AwsStorageProvider(IStorageProvider):
 
-    ARROW_NATIVE_FS_PROPERTY = "arrowNativeFs"
-    ARROW_NATIVE_FS_DEFAULT = False
-
     BUCKET_PROPERTY = "bucket"
     PREFIX_PROPERTY = "prefix"
     REGION_PROPERTY = "region"
@@ -56,6 +56,12 @@ class AwsStorageProvider(IStorageProvider):
     ACCESS_KEY_ID_PROPERTY = "accessKeyId"
     SECRET_ACCESS_KEY_PROPERTY = "secretAccessKey"
 
+    RUNTIME_FS_PROPERTY = "runtimeFs"
+    RUNTIME_FS_AUTO = "auto"
+    RUNTIME_FS_ARROW = "arrow"
+    RUNTIME_FS_BOTO3 = "boto3"
+    RUNTIME_FS_DEFAULT = RUNTIME_FS_AUTO
+
     ARROW_CLIENT_ARGS = {
         REGION_PROPERTY: "region",
         ENDPOINT_PROPERTY: "endpoint_override",
@@ -75,14 +81,25 @@ class AwsStorageProvider(IStorageProvider):
         self._log = _helpers.logger_for_object(self)
         self._properties = properties
 
-        self._arrow_native = _helpers.get_plugin_property_boolean(
-            properties, self.ARROW_NATIVE_FS_PROPERTY, self.ARROW_NATIVE_FS_DEFAULT)
+        self._runtime_fs = _helpers.get_plugin_property(
+            properties, self.RUNTIME_FS_PROPERTY) \
+            or self.RUNTIME_FS_DEFAULT
 
     def has_arrow_native(self) -> bool:
-        return True if self._arrow_native else False
+        if self._runtime_fs == self.RUNTIME_FS_ARROW:
+            return True
+        elif self._runtime_fs == self.RUNTIME_FS_AUTO:
+            return afs.S3FileSystem is not None
+        else:
+            return False
 
     def has_file_storage(self) -> bool:
-        return True if not self._arrow_native and boto_available else False
+        if self._runtime_fs == self.RUNTIME_FS_BOTO3:
+            return True
+        elif self._runtime_fs == self.RUNTIME_FS_AUTO:
+            return afs.S3FileSystem is None
+        else:
+            return False
 
     def get_arrow_native(self) -> afs.SubTreeFileSystem:
 
@@ -161,7 +178,8 @@ class AwsStorageProvider(IStorageProvider):
         raise ex.EStartup(message)
 
 
-plugins.PluginManager.register_plugin(IStorageProvider, AwsStorageProvider, ["S3"])
+if __aws_available:
+    plugins.PluginManager.register_plugin(IStorageProvider, AwsStorageProvider, ["S3"])
 
 
 # ----------------------------------------------------------------------------------------------------------------------
@@ -172,196 +190,194 @@ plugins.PluginManager.register_plugin(IStorageProvider, AwsStorageProvider, ["S3
 # It is likely to be removed in a future release
 
 
-if boto_available:
+class S3ObjectStorage(IFileStorage):
 
-    class S3ObjectStorage(IFileStorage):
+    # This is a quick implementation of IFileStorage on S3 using the boto3 AWS SDK
 
-        # This is a quick implementation of IFileStorage on S3 using the boto3 AWS SDK
+    def __init__(self, config: cfg.PluginConfig, client_args: dict):
 
-        def __init__(self, config: cfg.PluginConfig, client_args: dict):
-
-            self._log = _helpers.logger_for_object(self)
+        self._log = _helpers.logger_for_object(self)
 
-            self._properties = config.properties
-            self._bucket = _helpers.get_plugin_property(self._properties, AwsStorageProvider.BUCKET_PROPERTY)
-            self._prefix = _helpers.get_plugin_property(self._properties, AwsStorageProvider.PREFIX_PROPERTY) or ""
+        self._properties = config.properties
+        self._bucket = _helpers.get_plugin_property(self._properties, AwsStorageProvider.BUCKET_PROPERTY)
+        self._prefix = _helpers.get_plugin_property(self._properties, AwsStorageProvider.PREFIX_PROPERTY) or ""
 
-            if self._bucket is None or len(self._bucket.strip()) == 0:
-                message = f"Missing required config property [{AwsStorageProvider.BUCKET_PROPERTY}] for S3 storage"
-                self._log.error(message)
-                raise ex.EConfigParse(message)
+        if self._bucket is None or len(self._bucket.strip()) == 0:
+            message = f"Missing required config property [{AwsStorageProvider.BUCKET_PROPERTY}] for S3 storage"
+            self._log.error(message)
+            raise ex.EConfigParse(message)
 
-            self._client = boto3.client(**client_args)
+        self._client = boto3.client(**client_args)
 
-        def exists(self, storage_path: str) -> bool:
+    def exists(self, storage_path: str) -> bool:
 
-            try:
-                self._log.info(f"EXISTS [{storage_path}]")
+        try:
+            self._log.info(f"EXISTS [{storage_path}]")
 
-                object_key = self._resolve_path(storage_path)
-                self._client.head_object(Bucket=self._bucket, Key=object_key)
-                return True
+            object_key = self._resolve_path(storage_path)
+            self._client.head_object(Bucket=self._bucket, Key=object_key)
+            return True
 
-            except aws_ex.ClientError as error:
-                aws_code = error.response['Error']['Code']
-                if aws_code == str(http.HTTPStatus.NOT_FOUND.value):  # noqa
-                    return False
-                raise ex.EStorageRequest(f"Storage error: {str(error)}") from error
+        except aws_ex.ClientError as error:
+            aws_code = error.response['Error']['Code']
+            if aws_code == str(http.HTTPStatus.NOT_FOUND.value):  # noqa
+                return False
+            raise ex.EStorageRequest(f"Storage error: {str(error)}") from error
 
-        def size(self, storage_path: str) -> int:
+    def size(self, storage_path: str) -> int:
 
-            try:
-                self._log.info(f"SIZE [{storage_path}]")
+        try:
+            self._log.info(f"SIZE [{storage_path}]")
 
-                object_key = self._resolve_path(storage_path)
-                response = self._client.head_object(Bucket=self._bucket, Key=object_key)
-                return response['ContentLength']
+            object_key = self._resolve_path(storage_path)
+            response = self._client.head_object(Bucket=self._bucket, Key=object_key)
+            return response['ContentLength']
 
-            except aws_ex.ClientError as error:
-                raise ex.EStorageRequest(f"Storage error: {str(error)}") from error
+        except aws_ex.ClientError as error:
+            raise ex.EStorageRequest(f"Storage error: {str(error)}") from error
 
-        def stat(self, storage_path: str) -> FileStat:
+    def stat(self, storage_path: str) -> FileStat:
 
-            self._log.info(f"STAT [{storage_path}]")
+        self._log.info(f"STAT [{storage_path}]")
 
-            name = storage_path.split("/")[-1]
+        name = storage_path.split("/")[-1]
 
-            if self.exists(storage_path):
+        if self.exists(storage_path):
 
-                # Only OBJECTS can support stat atm
-                # Handling for directories needs to be changed, as part of refactor onto object storage
-                size = self.size(storage_path)
-                return FileStat(name, FileType.FILE, storage_path, size)
+            # Only OBJECTS can support stat atm
+            # Handling for directories needs to be changed, as part of refactor onto object storage
+            size = self.size(storage_path)
+            return FileStat(name, FileType.FILE, storage_path, size)
 
-            else:
+        else:
 
-                self.ls(storage_path)
-                return FileStat(name, FileType.DIRECTORY, storage_path, 0)
+            self.ls(storage_path)
+            return FileStat(name, FileType.DIRECTORY, storage_path, 0)
 
-        def ls(self, storage_path: str, recursive: bool = False) -> tp.List[FileStat]:
+    def ls(self, storage_path: str, recursive: bool = False) -> tp.List[FileStat]:
 
-            self._log.info(f"LS [{storage_path}]")
+        self._log.info(f"LS [{storage_path}]")
 
-            prefix = self._resolve_path(storage_path) + "/"
+        prefix = self._resolve_path(storage_path) + "/"
 
-            response = self._client.list_objects_v2(
-                Bucket=self._bucket,
-                Prefix=prefix,
-                Delimiter="/")
+        response = self._client.list_objects_v2(
+            Bucket=self._bucket,
+            Prefix=prefix,
+            Delimiter="/")
 
-            keys = []
+        keys = []
 
-            if "Contents" not in response and "CommonPrefixes" not in response:
-                raise ex.EStorageRequest(f"Storage prefix not found: [{storage_path}]")
+        if "Contents" not in response and "CommonPrefixes" not in response:
+            raise ex.EStorageRequest(f"Storage prefix not found: [{storage_path}]")
 
-            if "Contents" in response:
-                for entry in response["Contents"]:
-                    raw_key = entry["Key"]
-                    if raw_key == prefix:
-                        continue
-                    key = raw_key.replace(prefix, "")
-                    size = entry["Size"]
-                    mtime = entry["LastModified "]
-                    stat = FileStat(key, FileType.FILE, raw_key, size, mtime=mtime)
-                    keys.append(stat)
+        if "Contents" in response:
+            for entry in response["Contents"]:
+                raw_key = entry["Key"]
+                if raw_key == prefix:
+                    continue
+                key = raw_key.replace(prefix, "")
+                size = entry["Size"]
+                mtime = entry["LastModified "]
+                stat = FileStat(key, FileType.FILE, raw_key, size, mtime=mtime)
+                keys.append(stat)
 
-            if "CommonPrefixes" in response:
-                for raw_prefix in response["CommonPrefixes"]:
-                    common_prefix = raw_prefix.replace(prefix, "")
-                    stat = FileStat(common_prefix, FileType.DIRECTORY, raw_prefix, 0)
-                    keys.append(stat)
+        if "CommonPrefixes" in response:
+            for raw_prefix in response["CommonPrefixes"]:
+                common_prefix = raw_prefix.replace(prefix, "")
+                stat = FileStat(common_prefix, FileType.DIRECTORY, raw_prefix, 0)
+                keys.append(stat)
 
-            return keys
+        return keys
 
-        def mkdir(self, storage_path: str, recursive: bool = False):
+    def mkdir(self, storage_path: str, recursive: bool = False):
 
-            self._log.info(f"MKDIR [{storage_path}]")
+        self._log.info(f"MKDIR [{storage_path}]")
 
-            # No-op in object storage
-            pass
+        # No-op in object storage
+        pass
 
-        def rm(self, storage_path: str):
+    def rm(self, storage_path: str):
 
-            try:
-                self._log.info(f"RM [{storage_path}]")
+        try:
+            self._log.info(f"RM [{storage_path}]")
 
-                object_key = self._resolve_path(storage_path)
-                self._client.delete_object(Bucket=self._bucket, Key=object_key)
+            object_key = self._resolve_path(storage_path)
+            self._client.delete_object(Bucket=self._bucket, Key=object_key)
 
-            except aws_ex.ClientError as error:
-                raise ex.EStorageRequest(f"Storage error: {str(error)}") from error
+        except aws_ex.ClientError as error:
+            raise ex.EStorageRequest(f"Storage error: {str(error)}") from error
 
-        def rmdir(self, storage_path: str):
+    def rmdir(self, storage_path: str):
 
-            raise RuntimeError("RMDIR (recursive) not available for S3 storage")
+        raise RuntimeError("RMDIR (recursive) not available for S3 storage")
 
-        def read_bytes(self, storage_path: str) -> bytes:
+    def read_bytes(self, storage_path: str) -> bytes:
 
-            self._log.info(f"READ BYTES [{storage_path}]")
+        self._log.info(f"READ BYTES [{storage_path}]")
 
-            body = self._read_impl(storage_path)
-            return body.read()
+        body = self._read_impl(storage_path)
+        return body.read()
 
-        def read_byte_stream(self, storage_path: str) -> tp.BinaryIO:
+    def read_byte_stream(self, storage_path: str) -> tp.BinaryIO:
 
-            self._log.info(f"READ BYTE STREAM [{storage_path}]")
+        self._log.info(f"READ BYTE STREAM [{storage_path}]")
 
-            data = self.read_bytes(storage_path)
-            return io.BytesIO(data)
+        data = self.read_bytes(storage_path)
+        return io.BytesIO(data)
 
-        def _read_impl(self, storage_path: str) -> botocore.response.StreamingBody:
+    def _read_impl(self, storage_path: str):
 
-            try:
+        try:
 
-                object_key = self._resolve_path(storage_path)
-                response = self._client.get_object(Bucket=self._bucket, Key=object_key)
-                return response['Body']
+            object_key = self._resolve_path(storage_path)
+            response = self._client.get_object(Bucket=self._bucket, Key=object_key)
+            return response['Body']
 
-            except aws_ex.ClientError as error:
-                raise ex.EStorageRequest(f"Storage error: {str(error)}") from error
+        except aws_ex.ClientError as error:
+            raise ex.EStorageRequest(f"Storage error: {str(error)}") from error
 
-        def write_bytes(self, storage_path: str, data: bytes):
+    def write_bytes(self, storage_path: str, data: bytes):
 
-            try:
-                self._log.info(f"WRITE BYTES [{storage_path}]")
+        try:
+            self._log.info(f"WRITE BYTES [{storage_path}]")
 
-                object_key = self._resolve_path(storage_path)
+            object_key = self._resolve_path(storage_path)
 
-                self._client.put_object(
-                    Bucket=self._bucket,
-                    Key=object_key,
-                    Body=data)
+            self._client.put_object(
+                Bucket=self._bucket,
+                Key=object_key,
+                Body=data)
 
-            except aws_ex.ClientError as error:
-                raise ex.EStorageRequest(f"Storage error: {str(error)}") from error
+        except aws_ex.ClientError as error:
+            raise ex.EStorageRequest(f"Storage error: {str(error)}") from error
 
-        def write_byte_stream(self, storage_path: str) -> tp.BinaryIO:
+    def write_byte_stream(self, storage_path: str) -> tp.BinaryIO:
 
-            self._log.info(f"WRITE BYTE STREAM [{storage_path}]")
+        self._log.info(f"WRITE BYTE STREAM [{storage_path}]")
 
-            return self._AwsWriteBuf(self, storage_path)
+        return self._AwsWriteBuf(self, storage_path)
 
-        class _AwsWriteBuf(io.BytesIO):
+    class _AwsWriteBuf(io.BytesIO):
 
-            def __init__(self, storage, storage_path):
-                super().__init__()
-                self._storage = storage
-                self._storage_path = storage_path
-                self._written = False
+        def __init__(self, storage, storage_path):
+            super().__init__()
+            self._storage = storage
+            self._storage_path = storage_path
+            self._written = False
 
-            def close(self):
-                if not self._written:
-                    self.seek(0)
-                    data = self.read()
-                    self._storage.write_bytes(self._storage_path, data)
-                    self._written = True
+        def close(self):
+            if not self._written:
+                self.seek(0)
+                data = self.read()
+                self._storage.write_bytes(self._storage_path, data)
+                self._written = True
 
-        def _resolve_path(self, storage_path: str) -> str:
+    def _resolve_path(self, storage_path: str) -> str:
 
-            if self._prefix is None or self._prefix.strip() == "":
-                return storage_path
+        if self._prefix is None or self._prefix.strip() == "":
+            return storage_path
 
-            separator = "" if self._prefix.endswith("/") else "/"
-            full_path = self._prefix + separator + storage_path
+        separator = "" if self._prefix.endswith("/") else "/"
+        full_path = self._prefix + separator + storage_path
 
-            return full_path[1:] if full_path.startswith("/") else full_path
+        return full_path[1:] if full_path.startswith("/") else full_path

tracdap/rt/_plugins/storage_azure.py

@@ -0,0 +1,155 @@
+#  Copyright 2023 Accenture Global Solutions Limited
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import logging
+import typing as tp
+
+# TRAC interfaces
+import tracdap.rt.exceptions as ex
+import tracdap.rt.ext.plugins as plugins
+from tracdap.rt.ext.storage import *
+
+import pyarrow.fs as afs
+
+try:
+    # These dependencies are provided by the optional [azure] feature
+    # For local development, pip install -r requirements_plugins.txt
+    import azure.storage.blob as az_blob  # noqa
+    import adlfs  # noqa
+    __azure_available = True
+except ImportError:
+    adlfs = None
+    __azure_available = False
+
+# Set of common helpers across the core plugins (do not reference rt._impl)
+from . import _helpers
+
+
+class AzureBlobStorageProvider(IStorageProvider):
+
+    # This client depends on the Azure fsspec implementation, since there is no native implementation from Arrow
+    # To enable it, the tracdap package must be installed with the optional [azure] feature
+
+    # Current supported authentication mechanisms are "default" and "access_key"
+    # Client always uses location mode = primary, version aware = False
+
+    STORAGE_ACCOUNT_PROPERTY = "storageAccount"
+    CONTAINER_PROPERTY = "container"
+    PREFIX_PROPERTY = "prefix"
+
+    CREDENTIALS_PROPERTY = "credentials"
+    CREDENTIALS_DEFAULT = "default"
+    CREDENTIALS_ACCESS_KEY = "access_key"
+
+    ACCESS_KEY_PROPERTY = "accessKey"
+
+    RUNTIME_FS_PROPERTY = "runtimeFs"
+    RUNTIME_FS_AUTO = "auto"
+    RUNTIME_FS_FSSPEC = "fsspec"
+    RUNTIME_FS_DEFAULT = RUNTIME_FS_AUTO
+
+    def __init__(self, properties: tp.Dict[str, str]):
+
+        self._log = _helpers.logger_for_object(self)
+        self._properties = properties
+
+        self._runtime_fs = _helpers.get_plugin_property(
+            properties, self.RUNTIME_FS_PROPERTY) \
+            or self.RUNTIME_FS_DEFAULT
+
+        # The Azure SDK is very verbose with logging
+        # Avoid log noise by raising the log level for the Azure namespace
+        azure_log = _helpers.logger_for_namespace("azure.core")
+        azure_log.level = logging.WARNING
+
+    def has_arrow_native(self) -> bool:
+        return True
+
+    def get_arrow_native(self) -> afs.SubTreeFileSystem:
+
+        if self._runtime_fs == self.RUNTIME_FS_AUTO or self._runtime_fs == self.RUNTIME_FS_FSSPEC:
+            azure_fs = self.create_fsspec()
+        else:
+            message = f"Requested runtime FS [{self._runtime_fs}] is not available for Azure storage"
+            self._log.error(message)
+            raise ex.EStartup(message)
+
+        container = _helpers.get_plugin_property(self._properties, self.CONTAINER_PROPERTY)
+        prefix = _helpers.get_plugin_property(self._properties, self.PREFIX_PROPERTY)
+
+        if container is None or container.strip() == "":
+            message = f"Missing required config property [{self.CONTAINER_PROPERTY}] for Azure blob storage"
+            self._log.error(message)
+            raise ex.EConfigParse(message)
+
+        root_path = f"{container}/{prefix}" if prefix else container
+
+        return afs.SubTreeFileSystem(root_path, azure_fs)
+
+    def create_fsspec(self) -> afs.FileSystem:
+
+        azure_fsspec_args = self.setup_client_args()
+        azure_fsspec = adlfs.AzureBlobFileSystem(**azure_fsspec_args)
+
+        return afs.PyFileSystem(afs.FSSpecHandler(azure_fsspec))
+
+    def setup_client_args(self) -> tp.Dict[str, tp.Any]:
+
+        client_args = dict()
+
+        storage_account = _helpers.get_plugin_property(self._properties, self.STORAGE_ACCOUNT_PROPERTY)
+
+        if storage_account is None or len(storage_account.strip()) == 0:
+            message = f"Missing required config property [{self.STORAGE_ACCOUNT_PROPERTY}] for Azure blob storage"
+            self._log.error(message)
+            raise ex.EConfigParse(message)
+
+        client_args["account_name"] = storage_account
+
+        credentials = self.setup_credentials()
+        client_args.update(credentials)
+
+        return client_args
+
+    def setup_credentials(self):
+
+        # Only default (Google ADC) mechanism is supported
+        # Arrow GCP FS does also support access tokens, but ADC is probably all we ever need
+
+        mechanism = _helpers.get_plugin_property(self._properties, self.CREDENTIALS_PROPERTY)
+
+        if mechanism is None or len(mechanism) == 0 or mechanism.lower() == self.CREDENTIALS_DEFAULT:
+            self._log.info(f"Using [{self.CREDENTIALS_DEFAULT}] credentials mechanism")
+            return {"anon": False}
+
+        if mechanism == self.CREDENTIALS_ACCESS_KEY:
+
+            self._log.info(f"Using [{self.CREDENTIALS_ACCESS_KEY}] credentials mechanism")
+
+            access_key = _helpers.get_plugin_property(self._properties, self.ACCESS_KEY_PROPERTY)
+
+            if access_key is None or len(access_key.strip()) == 0:
+                message = f"Missing required config property [{self.ACCESS_KEY_PROPERTY}] for Azure blob storage"
+                raise ex.EConfigParse(message)
+
+            return {"account_key": access_key}
+
+        message = f"Unrecognised credentials mechanism: [{mechanism}]"
+        self._log.error(message)
+        raise ex.EStartup(message)
+
+
+# Only register the plugin if the [azure] feature is available
+if __azure_available:
+    plugins.PluginManager.register_plugin(IStorageProvider, AzureBlobStorageProvider, ["BLOB"])