Marmelator/pyinfra
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 36 Wiki Activity

refactoy: simplify storage connection logic

Browse Source
This commit is contained in:
Julius Unverfehrt 2024-01-25 09:08:51 +01:00
parent 8ff637d6ba
commit f6f56b8d8c
3 changed files with 64 additions and 69 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

121
pyinfra/storage/connection.py
View File

@ -16,83 +16,74 @@ from pyinfra.utils.cipher import decrypt
def get_storage(settings: Dynaconf, tenant_id: str = None) -> Storage:
"""Get storage connection based on settings.
If tenant_id is provided, gets storage connection information from tenant server instead.
The connections are cached based on the settings.cache_size value.
In the future, when the default storage from config is no longer needed (only multi-tenant storage will be used),
get_storage_from_tenant_id can replace this function directly.
"""Establishes a storage connection.
If tenant_id is provided, gets storage connection information from tenant server. These connections are cached.
Otherwise, gets storage connection information from settings.
"""
logger.info("Establishing storage connection...")
if tenant_id:
logger.info(f"Using tenant storage for {tenant_id}.")
return get_storage_from_tenant_id(tenant_id, settings)
else:
logger.info("Using default storage.")
return get_storage_from_settings(settings)
validate_settings(settings, multi_tenant_storage_validators)
return get_storage_for_tenant(
tenant_id,
settings.storage.tenant_server.endpoint,
settings.storage.tenant_server.public_key,
)
def get_storage_from_settings(settings: Dynaconf) -> Storage:
logger.info("Using default storage.")
validate_settings(settings, storage_validators)
@lru_cache(maxsize=settings.storage.cache_size)
def _get_storage(backend: str) -> Storage:
return storage_dispatcher[backend](settings)
return _get_storage(settings.storage.backend)
def get_storage_from_tenant_id(tenant_id: str, settings: Dynaconf) -> Storage:
validate_settings(settings, multi_tenant_storage_validators)
@lru_cache(maxsize=settings.storage.cache_size)
def _get_storage(tenant: str, endpoint: str, public_key: str) -> Storage:
response = requests.get(f"{endpoint}/{tenant}").json()
maybe_azure = response.get("azureStorageConnection")
maybe_s3 = response.get("s3StorageConnection")
assert (maybe_azure or maybe_s3) and not (maybe_azure and maybe_s3), "Only one storage backend can be used."
if maybe_azure:
connection_string = decrypt(public_key, maybe_azure["connectionString"])
backend = "azure"
storage_info = {
"storage": {
"azure": {
"connection_string": connection_string,
"container": maybe_azure["containerName"],
},
}
}
elif maybe_s3:
secret = decrypt(public_key, maybe_s3["secret"])
backend = "s3"
storage_info = {
"storage": {
"s3": {
"endpoint": maybe_s3["endpoint"],
"key": maybe_s3["key"],
"secret": secret,
"region": maybe_s3["region"],
"bucket": maybe_s3["bucketName"],
},
}
}
else:
raise Exception(f"Unknown storage backend in {response}.")
storage_settings = Dynaconf()
storage_settings.update(storage_info)
storage = storage_dispatcher[backend](storage_settings)
return storage
return _get_storage(tenant_id, settings.storage.tenant_server.endpoint, settings.storage.tenant_server.public_key)
return storage_dispatcher[settings.storage.backend](settings)
storage_dispatcher = {
"azure": get_azure_storage_from_settings,
"s3": get_s3_storage_from_settings,
}
@lru_cache(maxsize=10)
def get_storage_for_tenant(tenant: str, endpoint: str, public_key: str) -> Storage:
response = requests.get(f"{endpoint}/{tenant}").json()
maybe_azure = response.get("azureStorageConnection")
maybe_s3 = response.get("s3StorageConnection")
assert (maybe_azure or maybe_s3) and not (maybe_azure and maybe_s3), "Only one storage backend can be used."
if maybe_azure:
connection_string = decrypt(public_key, maybe_azure["connectionString"])
backend = "azure"
storage_info = {
"storage": {
"azure": {
"connection_string": connection_string,
"container": maybe_azure["containerName"],
},
}
}
elif maybe_s3:
secret = decrypt(public_key, maybe_s3["secret"])
backend = "s3"
storage_info = {
"storage": {
"s3": {
"endpoint": maybe_s3["endpoint"],
"key": maybe_s3["key"],
"secret": secret,
"region": maybe_s3["region"],
"bucket": maybe_s3["bucketName"],
},
}
}
else:
raise Exception(f"Unknown storage backend in {response}.")
storage_settings = Dynaconf()
storage_settings.update(storage_info)
storage = storage_dispatcher[backend](storage_settings)
return storage

4
tests/conftest.py
View File

@ -4,7 +4,7 @@ import pytest
from pyinfra.config.loader import load_settings, pyinfra_config_path
from pyinfra.queue.manager import QueueManager
from pyinfra.storage.connection import get_storage_from_settings
from pyinfra.storage.connection import get_storage
@pytest.fixture(scope="session")
@ -16,7 +16,7 @@ def settings():
def storage(storage_backend, settings):
settings.storage.backend = storage_backend
storage = get_storage_from_settings(settings)
storage = get_storage(settings)
storage.make_bucket()
yield storage

8
tests/unit_test/storage_test.py
View File

@ -5,7 +5,7 @@ from time import sleep
import pytest
from fastapi import FastAPI
from pyinfra.storage.connection import get_storage_from_tenant_id
from pyinfra.storage.connection import get_storage_for_tenant
from pyinfra.storage.utils import (
download_data_as_specified_in_message,
upload_data_as_specified_in_message,
@ -106,7 +106,11 @@ class TestMultiTenantStorage:
self, tenant_id, tenant_server_mock, settings, tenant_server_host, tenant_server_port
):
settings["storage"]["tenant_server"]["endpoint"] = f"http://{tenant_server_host}:{tenant_server_port}"
storage = get_storage_from_tenant_id(tenant_id, settings)
storage = get_storage_for_tenant(
tenant_id,
settings["storage"]["tenant_server"]["endpoint"],
settings["storage"]["tenant_server"]["public_key"],
)
storage.put_object("file", b"content")
data_received = storage.get_object("file")