lakehouse_engine.utils.storage.file_storage_functions

  1"""Module for common file storage functions."""
  2
  3import json
  4from abc import ABC
  5from typing import Any
  6from urllib.parse import ParseResult, urlparse
  7
  8import boto3
  9
 10from lakehouse_engine.utils.storage.dbfs_storage import DBFSStorage
 11from lakehouse_engine.utils.storage.local_fs_storage import LocalFSStorage
 12from lakehouse_engine.utils.storage.s3_storage import S3Storage
 13
 14
 15class FileStorageFunctions(ABC):  # noqa: B024
 16    """Class for common file storage functions."""
 17
 18    @classmethod
 19    def read_json(cls, path: str, disable_dbfs_retry: bool = False) -> Any:
 20        """Read a json file.
 21
 22        The file should be in a supported file system (e.g., s3, dbfs or
 23        local filesystem).
 24
 25        Args:
 26            path: path to the json file.
 27            disable_dbfs_retry: optional flag to disable file storage dbfs.
 28
 29        Returns:
 30            Dict with json file content.
 31        """
 32        url = urlparse(path, allow_fragments=False)
 33        if disable_dbfs_retry:
 34            return json.load(S3Storage.get_file_payload(url))
 35        elif url.scheme == "s3" and cls.is_boto3_configured():
 36            try:
 37                return json.load(S3Storage.get_file_payload(url))
 38            except Exception:
 39                return json.loads(DBFSStorage.get_file_payload(url))
 40        elif url.scheme == "file":
 41            return json.load(LocalFSStorage.get_file_payload(url))
 42        elif url.scheme in ["dbfs", "s3"]:
 43            return json.loads(DBFSStorage.get_file_payload(url))
 44        else:
 45            raise NotImplementedError(
 46                f"File storage protocol not implemented for {path}."
 47            )
 48
 49    @classmethod
 50    def read_sql(cls, path: str, disable_dbfs_retry: bool = False) -> Any:
 51        """Read a sql file.
 52
 53        The file should be in a supported file system (e.g., s3, dbfs or local
 54        filesystem).
 55
 56        Args:
 57            path: path to the sql file.
 58            disable_dbfs_retry: optional flag to disable file storage dbfs.
 59
 60        Returns:
 61            Content of the SQL file.
 62        """
 63        url = urlparse(path, allow_fragments=False)
 64        if disable_dbfs_retry:
 65            return S3Storage.get_file_payload(url).read().decode("utf-8")
 66        elif url.scheme == "s3" and cls.is_boto3_configured():
 67            try:
 68                return S3Storage.get_file_payload(url).read().decode("utf-8")
 69            except Exception:
 70                return DBFSStorage.get_file_payload(url)
 71        elif url.scheme == "file":
 72            return LocalFSStorage.get_file_payload(url).read()
 73        elif url.scheme in ["dbfs", "s3"]:
 74            return DBFSStorage.get_file_payload(url)
 75        else:
 76            raise NotImplementedError(
 77                f"Object storage protocol not implemented for {path}."
 78            )
 79
 80    @classmethod
 81    def write_payload(
 82        cls, path: str, url: ParseResult, content: str, disable_dbfs_retry: bool = False
 83    ) -> None:
 84        """Write payload into a file.
 85
 86        The file should be in a supported file system (e.g., s3, dbfs or local
 87        filesystem).
 88
 89        Args:
 90            path: path to validate the file type.
 91            url: url of the file.
 92            content: content to write into the file.
 93            disable_dbfs_retry: optional flag to disable file storage dbfs.
 94        """
 95        if disable_dbfs_retry:
 96            S3Storage.write_payload_to_file(url, content)
 97        elif path.startswith("s3://") and cls.is_boto3_configured():
 98            try:
 99                S3Storage.write_payload_to_file(url, content)
100            except Exception:
101                DBFSStorage.write_payload_to_file(url, content)
102        elif path.startswith(("s3://", "dbfs:/")):
103            DBFSStorage.write_payload_to_file(url, content)
104        else:
105            LocalFSStorage.write_payload_to_file(url, content)
106
107    @staticmethod
108    def is_boto3_configured() -> bool:
109        """Check if boto3 is able to locate credentials and properly configured.
110
111        If boto3 is not properly configured, we might want to try a different reader.
112        """
113        try:
114            boto3.client("sts").get_caller_identity()
115            return True
116        except Exception:
117            return False