--- /usr/local/etc/puppet/files/coreservers_package	2014-11-07 07:24:03.000000000 +0100
+++ /usr/local/etc/puppet/files/coreservers_package	2014-11-07 07:48:01.000000000 +0100
@@ -33,7 +33,7 @@
 
 service_been_upgraded()
   {
-    egrep '(Reinstalling (.*)\.\.\.|Upgrading (.*) from (.*) to (.*)\.\.\.)' ${UPGRADE_LOG} | grep ${SERVICE} | awk 'END{print NR}'
+    egrep '(Reinstalling (.*)\.\.\.|Upgrading (.*) from (.*) to (.*)\.\.\.)' ${UPGRADE_LOG} | egrep "${SERVICE}" | awk 'END{print NR}'
   }
 
 case "$1" in
@@ -56,11 +56,17 @@
 
   PKG_UPGRADE="pkg upgrade -y"
 
-  if [ ! -z $2 ]; then
-    if [ $2 = "-f" ]; then
-      PKG_UPGRADE="${PKG_UPGRADE} -f"
-    fi
-  fi
+  while [ "$#" -gt "0" ]; do
+    case "$1" in
+      "-f")
+        PKG_UPGRADE="${PKG_UPGRADE} -f"
+      ;;
+      *)
+        echo "Unknown argument '$1'" >&2
+      ;;
+    esac
+    shift
+  done
 
   ${PKG_UPGRADE} >> ${UPGRADE_LOG} 2>&1
   if [ "$?" -ne "0" ]; then
@@ -82,10 +88,22 @@
   fi
 
   for SERVICE in ${SERVICES}; do
+
+    if [ "${SERVICE}" = "mysql" ]; then
+      SERVICE="mysql|mariadb"
+    elif [ "${SERVICE}" = "puppet" ]; then
+      SERVICE="puppet|ruby"
+    fi
     if [ "$(service_been_upgraded)" -gt "0" ]; then
+      if [ "${SERVICE}" = "mysql|mariadb" ]; then
+        SERVICE="mysql-server"
+      elif [ "${SERVICE}" = "puppet|ruby" ]; then
+        SERVICE="puppet"
+      fi
       service ${SERVICE} restart >> ${UPGRADE_LOG} 2>&1
       echo "The service \"${SERVICE}\" has been restarted to load in an upgraded version." >> ${UPGRADE_LOG}
     fi
+
   done
 
   cp ${UPGRADE_LOG} /var/log/packageupgrade/packageupgrade.$(date +%Y%m%d-%H%M).log