Files
NotesAgentic/server sync/sync_server/storage.py
T

100 lines
3.5 KiB
Python

"""内容对象按 Vault 分区;测试磁盘适配器不作为生产对象存储。"""
from pathlib import Path
import hashlib
import os
import tempfile
import shutil
from contextlib import contextmanager
class DiskObjects:
def __init__(self, root: Path):
self.root = root
root.mkdir(parents=True, exist_ok=True)
def put(self, key: str, data: bytes):
target = self.root / key
target.parent.mkdir(parents=True, exist_ok=True)
fd, temp = tempfile.mkstemp(dir=target.parent)
try:
with os.fdopen(fd, "wb") as stream:
stream.write(data)
stream.flush()
os.fsync(stream.fileno())
os.replace(temp, target)
finally:
Path(temp).unlink(missing_ok=True)
def get(self, key: str) -> bytes:
return (self.root / key).read_bytes()
def put_file(self, key: str, path: Path, content_hash: str):
target = self.root / key
target.parent.mkdir(parents=True, exist_ok=True)
with tempfile.NamedTemporaryFile(dir=target.parent, delete=False) as stream:
temporary = Path(stream.name)
try:
with path.open("rb") as source:
shutil.copyfileobj(source, stream, 1024 * 1024)
stream.flush()
os.fsync(stream.fileno())
except BaseException:
stream.close()
temporary.unlink(missing_ok=True)
raise
try:
os.replace(temporary, target)
finally:
temporary.unlink(missing_ok=True)
@contextmanager
def open(self, key: str):
with (self.root / key).open("rb") as stream:
yield stream
def delete(self, key: str):
(self.root / key).unlink(missing_ok=True)
class S3Objects:
def __init__(self, endpoint: str, bucket: str):
import boto3
from botocore.config import Config
self.client = boto3.client("s3", endpoint_url=endpoint,
config=Config(connect_timeout=2, read_timeout=2,
retries={"max_attempts": 0}))
self.bucket = bucket
def put(self, key: str, data: bytes):
self.client.put_object(Bucket=self.bucket, Key=key, Body=data,
Metadata={"sha256": hashlib.sha256(data).hexdigest()})
def get(self, key: str) -> bytes:
response = self.client.get_object(Bucket=self.bucket, Key=key)
with response["Body"] as stream:
return stream.read()
def put_file(self, key: str, path: Path, content_hash: str):
with path.open("rb") as stream:
# Objects are capped at 100 MiB, well below S3's 5 GiB single-PUT
# limit. A direct streaming request has one explicit connection
# lifetime; constructing a transfer manager per completion can
# retain pooled MinIO connections under repeated multi-worker use.
self.client.put_object(
Bucket=self.bucket,
Key=key,
Body=stream,
ContentLength=path.stat().st_size,
Metadata={"sha256": content_hash},
)
@contextmanager
def open(self, key: str):
response = self.client.get_object(Bucket=self.bucket, Key=key)
with response["Body"] as stream:
yield stream
def delete(self, key: str):
self.client.delete_object(Bucket=self.bucket, Key=key)