| 1 | #include <IO/WriteBufferFromTemporaryFile.h> |
| 2 | #include <IO/ReadBufferFromFile.h> |
| 3 | |
| 4 | #include <fcntl.h> |
| 5 | |
| 6 | |
| 7 | namespace DB |
| 8 | { |
| 9 | |
| 10 | namespace ErrorCodes |
| 11 | { |
| 12 | extern const int CANNOT_OPEN_FILE; |
| 13 | extern const int CANNOT_SEEK_THROUGH_FILE; |
| 14 | } |
| 15 | |
| 16 | |
| 17 | WriteBufferFromTemporaryFile::WriteBufferFromTemporaryFile(std::unique_ptr<TemporaryFile> && tmp_file_) |
| 18 | : WriteBufferFromFile(tmp_file_->path(), DBMS_DEFAULT_BUFFER_SIZE, O_RDWR | O_TRUNC | O_CREAT, 0600), tmp_file(std::move(tmp_file_)) |
| 19 | {} |
| 20 | |
| 21 | |
| 22 | WriteBufferFromTemporaryFile::Ptr WriteBufferFromTemporaryFile::create(const std::string & tmp_dir) |
| 23 | { |
| 24 | return Ptr{new WriteBufferFromTemporaryFile(createTemporaryFile(tmp_dir))}; |
| 25 | } |
| 26 | |
| 27 | |
| 28 | class ReadBufferFromTemporaryWriteBuffer : public ReadBufferFromFile |
| 29 | { |
| 30 | public: |
| 31 | static ReadBufferPtr createFrom(WriteBufferFromTemporaryFile * origin) |
| 32 | { |
| 33 | int fd = origin->getFD(); |
| 34 | std::string file_name = origin->getFileName(); |
| 35 | |
| 36 | off_t res = lseek(fd, 0, SEEK_SET); |
| 37 | if (-1 == res) |
| 38 | throwFromErrnoWithPath("Cannot reread temporary file " + file_name, file_name, |
| 39 | ErrorCodes::CANNOT_SEEK_THROUGH_FILE); |
| 40 | |
| 41 | return std::make_shared<ReadBufferFromTemporaryWriteBuffer>(fd, file_name, std::move(origin->tmp_file)); |
| 42 | } |
| 43 | |
| 44 | ReadBufferFromTemporaryWriteBuffer(int fd_, const std::string & file_name_, std::unique_ptr<TemporaryFile> && tmp_file_) |
| 45 | : ReadBufferFromFile(fd_, file_name_), tmp_file(std::move(tmp_file_)) |
| 46 | {} |
| 47 | |
| 48 | std::unique_ptr<TemporaryFile> tmp_file; |
| 49 | }; |
| 50 | |
| 51 | |
| 52 | ReadBufferPtr WriteBufferFromTemporaryFile::getReadBufferImpl() |
| 53 | { |
| 54 | /// ignore buffer, write all data to file and reread it |
| 55 | next(); |
| 56 | |
| 57 | auto res = ReadBufferFromTemporaryWriteBuffer::createFrom(this); |
| 58 | |
| 59 | /// invalidate FD to avoid close(fd) in destructor |
| 60 | setFD(-1); |
| 61 | file_name = {}; |
| 62 | |
| 63 | return res; |
| 64 | } |
| 65 | |
| 66 | |
| 67 | WriteBufferFromTemporaryFile::~WriteBufferFromTemporaryFile() = default; |
| 68 | |
| 69 | |
| 70 | } |
| 71 | |