msintui/Intelligent_PID

1import os2import shutil3from azure.storage.blob import BlobServiceClient4from abc import ABC, abstractmethod5import json6 7class StorageInterface(ABC):8 9    @abstractmethod10    def save_file(self, file_path: str, content: bytes) -> str:11        pass12 13    @abstractmethod14    def load_file(self, file_path: str) -> bytes:15        pass16 17    @abstractmethod18    def list_files(self, directory: str) -> list[str]:19        pass20 21    @abstractmethod22    def file_exists(self, file_path: str) -> bool:23        pass24 25    @abstractmethod26    def delete_file(self, file_path: str) -> None:27        pass28 29    @abstractmethod30    def create_directory(self, directory: str) -> None:31        pass32 33    @abstractmethod34    def delete_directory(self, directory: str) -> None:35        pass36 37    @abstractmethod38    def upload(self, local_path: str, destination_path: str) -> None:39        pass40 41    @abstractmethod42    def append_file(self, file_path: str, content: bytes) -> None:43        pass44 45    @abstractmethod46    def get_modified_time(self, file_path: str) -> float:47        pass48 49    @abstractmethod50    def directory_exists(self, directory: str) -> bool:51        pass52 53    def load_json(self, file_path):54        """Load and parse JSON file."""55        try:56            with open(file_path, 'r', encoding='utf-8') as f:57                data = json.load(f)58            return data59        except Exception as e:60            print(f"Error loading JSON from {file_path}: {str(e)}")61            return None62 63 64class LocalStorage(StorageInterface):65 66    def save_file(self, file_path: str, content: bytes) -> str:67        os.makedirs(os.path.dirname(file_path), exist_ok=True)68        with open(file_path, 'wb') as f:69            f.write(content)70        return file_path71 72    def load_file(self, file_path: str) -> bytes:73        with open(file_path, 'rb') as f:74            return f.read()75 76    def list_files(self, directory: str) -> list[str]:77        return [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))]78 79    def file_exists(self, file_path: str) -> bool:80        return os.path.exists(file_path)81 82    def delete_file(self, file_path: str) -> None:83        os.remove(file_path)84 85    def create_directory(self, directory: str) -> None:86        os.makedirs(directory, exist_ok=True)87 88    def delete_directory(self, directory: str) -> None:89        shutil.rmtree(directory)90 91    def upload(self, local_path: str, destination_path: str) -> None:92        os.makedirs(os.path.dirname(destination_path), exist_ok=True)93        shutil.copy(local_path, destination_path)94 95    def append_file(self, file_path: str, content: bytes) -> None:96        os.makedirs(os.path.dirname(file_path), exist_ok=True)97        with open(file_path, 'ab') as f:98            f.write(content)99 100    def get_modified_time(self, file_path: str) -> float:101        return os.path.getmtime(file_path)102 103    def directory_exists(self, directory: str) -> bool:104        return self.file_exists(directory)105 106 107class BlobStorage(StorageInterface):108    """109        Writes to blob storage, using local disk as a cache110 111        TODO: Allow configuration of temp dir instead of just using the same paths in both local and remote112    """113    def __init__(self, connection_string: str, container_name: str):114        self.blob_service_client = BlobServiceClient.from_connection_string(connection_string)115        self.container_client = self.blob_service_client.get_container_client(container_name)116        self.local_storage = LocalStorage()117 118    def download(self, file_path: str) -> bytes:119        blob_client = self.container_client.get_blob_client(file_path)120        return blob_client.download_blob().readall()121 122    def sync(self, file_path: str) -> None:123        if not self.local_storage.file_exists(file_path):124            print(f"DEBUG: missing local version of {file_path} - downloading")125            self.local_storage.save_file(file_path, self.download(file_path))126        else:127            local_timestamp = self.local_storage.get_modified_time(file_path)128            remote_timestamp = self.get_modified_time(file_path)129            if local_timestamp < remote_timestamp:130                # We always write remotely before writing locally, so we expect local_timestamp to be > remote timestamp131                print(f"DBEUG: local version of {file_path} out of date - downloading")132                self.local_storage.save_file(file_path, self.download(file_path))133 134    135    def save_file(self, file_path: str, content: bytes) -> str:136        blob_client = self.container_client.get_blob_client(file_path)137        blob_client.upload_blob(content, overwrite=True)138        self.local_storage.save_file(file_path, content)139        return file_path140 141    def load_file(self, file_path: str) -> bytes:142        self.sync(file_path)143        return self.local_storage.load_file(file_path)144 145    def list_files(self, directory: str) -> list[str]:146        return [blob.name for blob in self.container_client.list_blobs(name_starts_with=directory)]147 148    def file_exists(self, file_path: str) -> bool:149        blob_client = self.container_client.get_blob_client(file_path)150        return blob_client.exists()151 152    def delete_file(self, file_path: str) -> None:153        self.local_storage.delete_file(file_path)154        blob_client = self.container_client.get_blob_client(file_path)155        blob_client.delete_blob()156 157    def create_directory(self, directory: str) -> None:158        # Blob storage doesn't have directories, so only create it locally159        self.local_storage.create_directory(directory)160 161    def delete_directory(self, directory: str) -> None:162        self.local_storage.delete_directory(directory)163        blobs_to_delete = self.container_client.list_blobs(name_starts_with=directory)164        for blob in blobs_to_delete:165            self.container_client.delete_blob(blob.name)166 167    def upload(self, local_path: str, destination_path: str) -> None:168        with open(local_path, "rb") as data:169            blob_client = self.container_client.get_blob_client(destination_path)170            blob_client.upload_blob(data, overwrite=True)171        self.local_storage.upload(local_path, destination_path)172 173    def append_file(self, file_path: str, content: bytes) -> None:174        blob_client = self.container_client.get_blob_client(file_path)175        if not blob_client.exists():176            blob_client.create_append_blob()177        else:178            self.sync(file_path)179 180        blob_client.append_block(content)181        self.local_storage.append_file(file_path, content)182 183    def get_modified_time(self, file_path: str) -> float:184        blob_client = self.container_client.get_blob_client(file_path)185        properties = blob_client.get_blob_properties()186        # Convert the UTC datetime to a UNIX timestamp187        return properties.last_modified.timestamp()188 189    def directory_exists(self, directory: str) -> bool:190        blobs = self.container_client.list_blobs(name_starts_with=directory)191        return next(blobs, None) is not None192 193 194class StorageFactory:195    @staticmethod196    def get_storage() -> StorageInterface:197        storage_type = os.getenv('STORAGE_TYPE', 'local').lower()198        if storage_type == 'local':199            return LocalStorage()200        elif storage_type == 'blob':201            connection_string = os.getenv('AZURE_STORAGE_CONNECTION_STRING')202            container_name = os.getenv('AZURE_STORAGE_CONTAINER_NAME')203            if not connection_string or not container_name:204                raise ValueError("Azure Blob Storage connection string and container name must be set")205            return BlobStorage(connection_string, container_name)206        else:207            raise ValueError(f"Unsupported storage type: {storage_type}")208 209