mongo-ruby-driver/lib/mongo/gridfs/grid_io.rb

247 lines
7.6 KiB
Ruby
Raw Normal View History

2010-02-12 23:03:07 +00:00
# --
# Copyright (C) 2008-2009 10gen Inc.
#
# 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.
# ++
module Mongo
2010-02-18 21:31:25 +00:00
# WARNING: This is part of a new, experimental GridFS API. Subject to change.
2010-02-12 23:03:07 +00:00
class GridIO
DEFAULT_CHUNK_SIZE = 256 * 1024
2010-02-18 21:31:25 +00:00
DEFAULT_CONTENT_TYPE = 'binary/octet-stream'
2010-02-12 23:03:07 +00:00
2010-02-18 21:31:25 +00:00
attr_reader :content_type, :chunk_size, :upload_date, :files_id, :filename, :metadata
2010-02-12 23:03:07 +00:00
2010-02-18 21:31:25 +00:00
def initialize(files, chunks, filename, mode, filesystem, opts={})
2010-02-12 23:03:07 +00:00
@files = files
@chunks = chunks
@filename = filename
@mode = mode
@content_type = opts[:content_type] || DEFAULT_CONTENT_TYPE
@chunk_size = opts[:chunk_size] || DEFAULT_CHUNK_SIZE
2010-02-18 21:31:25 +00:00
@files_id = opts[:_id]
2010-02-12 23:03:07 +00:00
2010-02-18 21:31:25 +00:00
case @mode
when 'r' then init_read(filesystem, opts)
when 'w' then init_write(opts)
else
raise GridError, "Invalid file mode #{@mode}. Valid options include 'r' and 'w'."
end
2010-02-12 23:03:07 +00:00
end
# Read the data from the file. If a length if specified, will read from the
# current file position.
#
# @param [Integer] length
#
# @return [String]
# the data in the file
def read(length=nil)
return '' if length == 0
return read_all if length.nil? && @file_position.zero?
buf = ''
while true
buf << @current_chunk['data'].to_s[@chunk_position..-1]
if buf.length >= length
return buf[0...length]
else
@current_chunk = get_chunk(@current_chunk['n'] + 1)
end
end
buf
end
2010-02-18 21:31:25 +00:00
alias :data :read
2010-02-12 23:03:07 +00:00
# Write the given string (binary) data to the file.
#
# @param [String] string
# the data to write
#
# @return [Integer]
# the number of bytes written.
def write(string)
raise GridError, "#{@filename} not opened for write" unless @mode[0] == ?w
# Since Ruby 1.9.1 doesn't necessarily store one character per byte.
if string.respond_to?(:force_encoding)
string.force_encoding("binary")
end
to_write = string.length
while (to_write > 0) do
if @current_chunk && @chunk_position == @chunk_size
next_chunk_number = @current_chunk['n'] + 1
@current_chunk = create_chunk(next_chunk_number)
end
chunk_available = @chunk_size - @chunk_position
step_size = (to_write > chunk_available) ? chunk_available : to_write
2010-02-18 21:31:25 +00:00
@current_chunk['data'] = Binary.new((@current_chunk['data'].to_s << string[-to_write, step_size]).unpack("c*"))
2010-02-12 23:03:07 +00:00
@chunk_position += step_size
to_write -= step_size
save_chunk(@current_chunk)
end
string.length - to_write
end
# Position the file pointer at the provided location.
#
# @param [Integer] pos
# the number of bytes to advance the file pointer. this can be a negative
# number.
# @param [Integer] whence
# one of IO::SEEK_CUR, IO::SEEK_END, or IO::SEEK_SET
#
# @return [Integer] the new file position
def seek(pos, whence=IO::SEEK_SET)
raise GridError, "Seek is only allowed in read mode." unless @mode == 'r'
target_pos = case whence
when IO::SEEK_CUR
@file_position + pos
when IO::SEEK_END
@file_length + pos
when IO::SEEK_SET
pos
end
new_chunk_number = (target_pos / @chunk_size).to_i
if new_chunk_number != @current_chunk['n']
save_chunk(@current_chunk) if @mode[0] == ?w
@current_chunk = get_chunk(new_chunk_number)
end
@file_position = target_pos
@chunk_position = @file_position % @chunk_size
@file_position
end
# The current position of the file.
#
# @return [Integer]
def tell
@file_position
end
# Creates or updates the document storing the chunks' metadata
# in the files collection. The file exists only after this method
# is called.
#
# This method will be invoked automatically
# on GridIO#open. Otherwise, it must be called manually.
#
# @return [True]
def close
if @mode[0] == ?w
2010-02-18 21:31:25 +00:00
@upload_date = Time.now.utc
2010-02-12 23:03:07 +00:00
@files.insert(to_mongo_object)
end
true
end
2010-02-18 21:31:25 +00:00
def inspect
"_id: #{@files_id}"
end
2010-02-12 23:03:07 +00:00
private
def create_chunk(n)
chunk = OrderedHash.new
chunk['_id'] = Mongo::ObjectID.new
chunk['n'] = n
chunk['files_id'] = @files_id
chunk['data'] = ''
2010-02-17 21:48:23 +00:00
@chunk_position = 0
2010-02-12 23:03:07 +00:00
chunk
end
2010-02-17 21:48:23 +00:00
# TODO: Perhaps use an upsert here instead?
2010-02-12 23:03:07 +00:00
def save_chunk(chunk)
@chunks.remove('_id' => chunk['_id'])
@chunks.insert(chunk)
end
def get_chunk(n)
chunk = @chunks.find({'files_id' => @files_id, 'n' => n}).next_document
@chunk_position = 0
chunk || {}
end
def last_chunk_number
(@file_length / @chunk_size).to_i
end
# An optimized read method for reading the whole file.
def read_all
buf = ''
while true
buf << @current_chunk['data'].to_s
break if @current_chunk['n'] == last_chunk_number
@current_chunk = get_chunk(@current_chunk['n'] + 1)
end
buf
end
# Initialize based on whether the supplied file exists.
2010-02-18 21:31:25 +00:00
def init_read(filesystem, opts)
if filesystem
doc = @files.find({'filename' => @filename}, :sort => [["uploadDate", -1]], :limit => 1).next_document
raise GridError, "Could not open file with filename #{@filename}" unless doc
2010-02-12 23:03:07 +00:00
else
2010-02-18 21:31:25 +00:00
doc = @files.find({'_id' => @files_id}).next_document
raise GridError, "Could not open file with id #{@files_id}" unless doc
2010-02-12 23:03:07 +00:00
end
2010-02-18 21:31:25 +00:00
@files_id = doc['_id']
@content_type = doc['contentType']
@chunk_size = doc['chunkSize']
@upload_date = doc['uploadDate']
@aliases = doc['aliases']
@file_length = doc['length']
@metadata = doc['metadata']
@md5 = doc['md5']
@filename = doc['filename']
@current_chunk = get_chunk(0)
@file_position = 0
2010-02-12 23:03:07 +00:00
end
# Validates and sets up the class for the given file mode.
2010-02-18 21:31:25 +00:00
def init_write(opts)
@files_id = opts[:_id] || Mongo::ObjectID.new
@content_type = opts[:content_type] || @content_type || DEFAULT_CONTENT_TYPE
@chunk_size = opts[:chunk_size] || @chunk_size || DEFAULT_CHUNK_SIZE
@file_length = 0
@metadata = opts[:metadata] if opts[:metadata]
@current_chunk = create_chunk(0)
@file_position = 0
2010-02-12 23:03:07 +00:00
end
def to_mongo_object
h = OrderedHash.new
h['_id'] = @files_id
h['filename'] = @filename
h['contentType'] = @content_type
h['length'] = @current_chunk ? @current_chunk['n'] * @chunk_size + @chunk_position : 0
h['chunkSize'] = @chunk_size
h['uploadDate'] = @upload_date
h['aliases'] = @aliases
h['metadata'] = @metadata
# Get a server-side md5.
md5_command = OrderedHash.new
md5_command['filemd5'] = @files_id
md5_command['root'] = 'fs'
h['md5'] = @files.db.command(md5_command)['md5']
h
end
end
end