forked-synapse/synapse/rest/media/v1/media_storage.py

266 lines
8.6 KiB
Python
Raw Normal View History

2018-01-08 12:07:30 -05:00
# -*- coding: utf-8 -*-
# Copyright 2018 New Vecotr Ltd
#
# 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.
2018-07-09 02:09:20 -04:00
import contextlib
import logging
import os
import shutil
import sys
2018-01-08 12:07:30 -05:00
import six
from twisted.internet import defer
2018-07-09 02:09:20 -04:00
from twisted.protocols.basic import FileSender
2018-01-08 12:07:30 -05:00
from synapse.logging.context import defer_to_thread, make_deferred_yieldable
2018-01-17 11:56:23 -05:00
from synapse.util.file_consumer import BackgroundFileConsumer
2018-01-08 12:07:30 -05:00
2018-07-09 02:09:20 -04:00
from ._base import Responder
2018-01-17 11:56:23 -05:00
2018-01-08 12:07:30 -05:00
logger = logging.getLogger(__name__)
class MediaStorage(object):
"""Responsible for storing/fetching files from local sources.
2018-01-12 06:15:31 -05:00
Args:
hs (synapse.server.Homeserver)
2018-01-12 06:15:31 -05:00
local_media_directory (str): Base path where we store media on disk
filepaths (MediaFilePaths)
storage_providers ([StorageProvider]): List of StorageProvider that are
used to fetch and store files.
2018-01-08 12:07:30 -05:00
"""
def __init__(self, hs, local_media_directory, filepaths, storage_providers):
self.hs = hs
2018-01-08 12:07:30 -05:00
self.local_media_directory = local_media_directory
self.filepaths = filepaths
2018-01-08 12:19:55 -05:00
self.storage_providers = storage_providers
2018-01-08 12:07:30 -05:00
@defer.inlineCallbacks
def store_file(self, source, file_info):
"""Write `source` to the on disk media store, and also any other
configured storage providers
Args:
source: A file like object that should be written
file_info (FileInfo): Info about the file to store
Returns:
Deferred[str]: the file path written to in the primary media store
"""
2018-02-14 12:55:18 -05:00
with self.store_into_file(file_info) as (f, fname, finish_cb):
# Write to the main repository
yield defer_to_thread(
2019-06-20 05:32:02 -04:00
self.hs.get_reactor(), _write_file_synchronously, source, f
)
2018-02-14 12:55:18 -05:00
yield finish_cb()
return fname
2018-01-08 12:07:30 -05:00
@contextlib.contextmanager
def store_into_file(self, file_info):
"""Context manager used to get a file like object to write into, as
described by file_info.
2018-01-12 06:15:31 -05:00
Actually yields a 3-tuple (file, fname, finish_cb), where file is a file
like object that can be written to, fname is the absolute path of file
on disk, and finish_cb is a function that returns a Deferred.
fname can be used to read the contents from after upload, e.g. to
generate thumbnails.
finish_cb must be called and waited on after the file has been
successfully been written to. Should not be called if there was an
error.
2018-01-08 12:07:30 -05:00
Args:
file_info (FileInfo): Info about the file to store
Example:
with media_storage.store_into_file(info) as (f, fname, finish_cb):
# .. write into f ...
yield finish_cb()
"""
path = self._file_info_to_path(file_info)
fname = os.path.join(self.local_media_directory, path)
dirname = os.path.dirname(fname)
if not os.path.exists(dirname):
os.makedirs(dirname)
finished_called = [False]
2018-01-08 12:19:55 -05:00
@defer.inlineCallbacks
2018-01-08 12:07:30 -05:00
def finish():
2018-01-08 12:19:55 -05:00
for provider in self.storage_providers:
yield provider.store_file(path, file_info)
2018-01-08 12:07:30 -05:00
finished_called[0] = True
try:
with open(fname, "wb") as f:
yield f, fname, finish
2018-01-12 09:44:02 -05:00
except Exception:
t, v, tb = sys.exc_info()
2018-01-08 12:07:30 -05:00
try:
os.remove(fname)
except Exception:
pass
six.reraise(t, v, tb)
2018-01-08 12:07:30 -05:00
if not finished_called:
2018-01-12 06:15:31 -05:00
raise Exception("Finished callback not called")
2018-01-08 12:07:30 -05:00
@defer.inlineCallbacks
def fetch_media(self, file_info):
"""Attempts to fetch media described by file_info from the local cache
and configured storage providers.
Args:
file_info (FileInfo)
Returns:
2018-01-12 06:15:31 -05:00
Deferred[Responder|None]: Returns a Responder if the file was found,
2018-01-08 12:07:30 -05:00
otherwise None.
"""
path = self._file_info_to_path(file_info)
local_path = os.path.join(self.local_media_directory, path)
if os.path.exists(local_path):
return FileResponder(open(local_path, "rb"))
2018-01-08 12:07:30 -05:00
2018-01-08 12:19:55 -05:00
for provider in self.storage_providers:
res = yield provider.fetch(path, file_info)
if res:
return res
2018-01-08 12:19:55 -05:00
return None
2018-01-08 12:07:30 -05:00
@defer.inlineCallbacks
def ensure_media_is_in_local_cache(self, file_info):
"""Ensures that the given file is in the local cache. Attempts to
download it from storage providers if it isn't.
Args:
file_info (FileInfo)
Returns:
Deferred[str]: Full path to local file
"""
path = self._file_info_to_path(file_info)
local_path = os.path.join(self.local_media_directory, path)
if os.path.exists(local_path):
return local_path
2018-01-17 11:56:23 -05:00
dirname = os.path.dirname(local_path)
if not os.path.exists(dirname):
os.makedirs(dirname)
for provider in self.storage_providers:
res = yield provider.fetch(path, file_info)
if res:
with res:
consumer = BackgroundFileConsumer(
2019-06-20 05:32:02 -04:00
open(local_path, "wb"), self.hs.get_reactor()
)
2018-01-17 11:56:23 -05:00
yield res.write_to_consumer(consumer)
yield consumer.wait()
return local_path
raise Exception("file could not be found")
2018-01-08 12:07:30 -05:00
def _file_info_to_path(self, file_info):
"""Converts file_info into a relative path.
2018-01-12 06:15:31 -05:00
The path is suitable for storing files under a directory, e.g. used to
store files on local FS under the base media repository directory.
Args:
file_info (FileInfo)
Returns:
str
2018-01-08 12:07:30 -05:00
"""
if file_info.url_cache:
if file_info.thumbnail:
return self.filepaths.url_cache_thumbnail_rel(
media_id=file_info.file_id,
width=file_info.thumbnail_width,
height=file_info.thumbnail_height,
content_type=file_info.thumbnail_type,
method=file_info.thumbnail_method,
)
2018-01-08 12:07:30 -05:00
return self.filepaths.url_cache_filepath_rel(file_info.file_id)
if file_info.server_name:
if file_info.thumbnail:
return self.filepaths.remote_media_thumbnail_rel(
server_name=file_info.server_name,
file_id=file_info.file_id,
width=file_info.thumbnail_width,
height=file_info.thumbnail_height,
content_type=file_info.thumbnail_type,
2019-06-20 05:32:02 -04:00
method=file_info.thumbnail_method,
2018-01-08 12:07:30 -05:00
)
return self.filepaths.remote_media_filepath_rel(
2019-06-20 05:32:02 -04:00
file_info.server_name, file_info.file_id
2018-01-08 12:07:30 -05:00
)
if file_info.thumbnail:
return self.filepaths.local_media_thumbnail_rel(
media_id=file_info.file_id,
width=file_info.thumbnail_width,
height=file_info.thumbnail_height,
content_type=file_info.thumbnail_type,
2019-06-20 05:32:02 -04:00
method=file_info.thumbnail_method,
2018-01-08 12:07:30 -05:00
)
2019-06-20 05:32:02 -04:00
return self.filepaths.local_media_filepath_rel(file_info.file_id)
2018-01-08 12:07:30 -05:00
2018-02-14 12:55:18 -05:00
def _write_file_synchronously(source, dest):
"""Write `source` to the file like `dest` synchronously. Should be called
2018-01-08 12:07:30 -05:00
from a thread.
Args:
2018-02-14 12:55:18 -05:00
source: A file like object that's to be written
dest: A file like object to be written to
2018-01-08 12:07:30 -05:00
"""
source.seek(0) # Ensure we read from the start of the file
2018-02-14 12:55:18 -05:00
shutil.copyfileobj(source, dest)
2018-01-08 12:07:30 -05:00
class FileResponder(Responder):
"""Wraps an open file that can be sent to a request.
Args:
open_file (file): A file like object to be streamed ot the client,
is closed when finished streaming.
"""
2019-06-20 05:32:02 -04:00
2018-01-08 12:07:30 -05:00
def __init__(self, open_file):
self.open_file = open_file
def write_to_consumer(self, consumer):
return make_deferred_yieldable(
FileSender().beginFileTransfer(self.open_file, consumer)
)
2018-01-08 12:07:30 -05:00
2018-01-12 08:32:03 -05:00
def __exit__(self, exc_type, exc_val, exc_tb):
2018-01-08 12:07:30 -05:00
self.open_file.close()