From ff63165bdd7b9930eed57fd86adeb3b7dcd1d6b4 Mon Sep 17 00:00:00 2001
From: Kyle Banker <kylebanker@gmail.com>
Date: Tue, 23 Feb 2010 18:03:11 -0500
Subject: [PATCH] content type detection for gridfs

---
 lib/mongo/gridfs/grid.rb    |  2 +-
 lib/mongo/gridfs/grid_io.rb | 16 +++++++++++++---
 test/grid_io_test.rb        | 21 ++++++++++++++++++++-
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/lib/mongo/gridfs/grid.rb b/lib/mongo/gridfs/grid.rb
index bc49acf..8384951 100644
--- a/lib/mongo/gridfs/grid.rb
+++ b/lib/mongo/gridfs/grid.rb
@@ -16,7 +16,7 @@
 
 module Mongo
 
-  # Implements the basic MongoDB GridFS specification.
+  # Implementation of the MongoDB GridFS specification. A file store.
   class Grid
     DEFAULT_FS_NAME = 'fs'
 
diff --git a/lib/mongo/gridfs/grid_io.rb b/lib/mongo/gridfs/grid_io.rb
index ea25ace..29f9137 100644
--- a/lib/mongo/gridfs/grid_io.rb
+++ b/lib/mongo/gridfs/grid_io.rb
@@ -15,6 +15,7 @@
 # ++
 
 require 'digest/md5'
+require 'mime/types'
 
 module Mongo
 
@@ -283,8 +284,8 @@ module Mongo
     # Initialize the class for writing a file.
     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
+      @content_type  = opts[:content_type] || get_content_type || DEFAULT_CONTENT_TYPE
+      @chunk_size    = opts[:chunk_size] || DEFAULT_CHUNK_SIZE
       @file_length   = 0
       @metadata      = opts[:metadata] if opts[:metadata]
 
@@ -315,11 +316,20 @@ module Mongo
       if @safe
         @client_md5 = @local_md5.hexdigest
         if @local_md5 != @server_md5
-          raise @local_md5 != @server_md5GridError, "File on server failed MD5 check"
+          raise @local_md5 != @server_md5, "File on server failed MD5 check"
         end
       else
         @server_md5
       end
     end
+
+    # Determine the content type based on the filename.
+    def get_content_type
+      if @filename
+        if types = MIME::Types.type_for(@filename)
+          types.first.simplified unless types.empty?
+        end
+      end
+    end
   end
 end
diff --git a/test/grid_io_test.rb b/test/grid_io_test.rb
index ad307a7..b93ee63 100644
--- a/test/grid_io_test.rb
+++ b/test/grid_io_test.rb
@@ -35,7 +35,6 @@ class GridIOTest < Test::Unit::TestCase
     end
 
     context "Grid MD5 check" do
-
       should "run in safe mode" do
         file = GridIO.new(@files, @chunks, 'smallfile', 'w', :safe => true)
         file.write("DATA" * 100)
@@ -51,6 +50,26 @@ class GridIOTest < Test::Unit::TestCase
         assert_equal file.server_md5, file.client_md5
       end
     end
+
+    context "Content types" do
+      should "determine common content types from the extension" do
+        file = GridIO.new(@files, @chunks, 'sample.pdf', 'w')
+        assert_equal 'application/pdf', file.content_type
+
+        file = GridIO.new(@files, @chunks, 'sample.txt', 'w')
+        assert_equal 'text/plain', file.content_type
+      end
+
+      should "default to binary/octet-stream when type is unknown" do
+        file = GridIO.new(@files, @chunks, 'sample.l33t', 'w')
+        assert_equal 'binary/octet-stream', file.content_type
+      end
+
+      should "use any provided content type by default" do
+        file = GridIO.new(@files, @chunks, 'sample.l33t', 'w', :content_type => 'image/jpg')
+        assert_equal 'image/jpg', file.content_type
+      end
+    end
   end
 
 end