allow to use SSL and fix console scripts invocation

mongodbtools/collection_stats.py

@@ -7,9 +7,18 @@
 
 from prettytable import PrettyTable
 import psutil
-from pymongo import Connection
+from socket import getfqdn
 from pymongo import ReadPreference
 from optparse import OptionParser
+from distutils.version import StrictVersion
+import pymongo
+
+HAS_PYMONGO3 = bool(StrictVersion(pymongo.version) >= StrictVersion('3.0'))
+
+if HAS_PYMONGO3:
+    from pymongo import MongoClient
+else:
+    from pymongo import Connection as MongoClient  # pylint: disable=E0611
 
 def compute_signature(index):
     signature = index["ns"]
@@ -18,7 +27,7 @@ def compute_signature(index):
     return signature
 
 def get_collection_stats(database, collection):
-    print "Checking DB: %s" % collection.full_name
+    print("Checking DB: %s" % collection.full_name)
     return database.command("collstats", collection.name)
 
 def get_cli_options():
@@ -50,18 +59,34 @@ def get_cli_options():
                       default="",
                       metavar="PASSWORD",
                       help="Admin password if authentication is enabled")
+    parser.add_option("--ssl-cert",
+                      dest="ssl_certfile",
+                      default=None,
+                      metavar="CERTIFICATE",
+                      help="SSL Certificate to use is SSL is enabled (only with pymongo >= 3)")
+    parser.add_option("--ssl-ca-certs",
+                      dest="ssl_ca_certs",
+                      default=None,
+                      metavar="CA",
+                      help="SSL Certificate of CA for certificate validation if SSL is enabled (only with pymongo >= 3)")
 
     (options, args) = parser.parse_args()
 
     return options
 
-def get_connection(host, port, username, password):
+def get_connection(host, port, username, password, ssl_certfile=None, ssl_ca_certs=None):
     userPass = ""
     if username and password:
         userPass = username + ":" + password + "@"
 
     mongoURI = "mongodb://" + userPass + host + ":" + str(port)
-    return Connection(host=mongoURI, read_preference=ReadPreference.SECONDARY)
+
+    conn_kwargs = dict(host=mongoURI, read_preference=ReadPreference.SECONDARY)
+
+    if HAS_PYMONGO3:
+        conn_kwargs.update(dict(ssl_certfile=ssl_certfile, ssl_ca_certs=ssl_ca_certs))
+
+    return MongoClient(**conn_kwargs)
 
 # From http://www.5dollarwhitebox.org/drupal/node/84
 def convert_bytes(bytes):
@@ -83,7 +108,10 @@ def convert_bytes(bytes):
         size = '%.2fb' % bytes
     return size
 
-def main(options):
+def main(options=None):
+    if options is None:
+        options = get_cli_options()
+
     summary_stats = {
         "count" : 0,
         "size" : 0,
@@ -92,7 +120,8 @@ def main(options):
     }
     all_stats = []
 
-    connection = get_connection(options.host, options.port, options.user, options.password)
+    connection = get_connection(options.host, options.port, options.user, options.password,
+                                options.ssl_certfile, options.ssl_ca_certs)
 
     all_db_stats = {}
 
@@ -145,18 +174,18 @@ def main(options):
                        ])
 
     print
-    print x.get_string(sortby="% Size")
-    print "Total Documents:", summary_stats["count"]
-    print "Total Data Size:", convert_bytes(summary_stats["size"])
-    print "Total Index Size:", convert_bytes(summary_stats["indexSize"])
-    print "Total Storage Size:", convert_bytes(summary_stats["storageSize"])
+    print(x.get_string(sortby="% Size"))
+    print("Total Documents: %s" % summary_stats["count"])
+    print("Total Data Size: %s" % convert_bytes(summary_stats["size"]))
+    print("Total Index Size: %s" % convert_bytes(summary_stats["indexSize"]))
+    print("Total Storage Size: %s" % convert_bytes(summary_stats["storageSize"]))
 
     # this is only meaningful if we're running the script on localhost
-    if options.host == "localhost":
+    if options.host == "localhost" or options.host == getfqdn():
         ram_headroom = psutil.phymem_usage()[0] - summary_stats["indexSize"]
-        print "RAM Headroom:", convert_bytes(ram_headroom)
-        print "RAM Used: %s (%s%%)" % (convert_bytes(psutil.phymem_usage()[1]), psutil.phymem_usage()[3])
-        print "Available RAM Headroom:", convert_bytes((100 - psutil.phymem_usage()[3]) / 100 * ram_headroom)
+        print("RAM Headroom: %s" % convert_bytes(ram_headroom))
+        print("RAM Used: %s (%s%%)" % (convert_bytes(psutil.phymem_usage()[1]), psutil.phymem_usage()[3]))
+        print("Available RAM Headroom: %s" % convert_bytes((100 - psutil.phymem_usage()[3]) / 100 * ram_headroom))
 
 if __name__ == "__main__":
     options = get_cli_options()

mongodbtools/index_stats.py

@@ -7,9 +7,18 @@
 
 from prettytable import PrettyTable
 import psutil
-from pymongo import Connection
+from socket import getfqdn
 from pymongo import ReadPreference
 from optparse import OptionParser
+from distutils.version import StrictVersion
+import pymongo
+
+HAS_PYMONGO3 = bool(StrictVersion(pymongo.version) >= StrictVersion('3.0'))
+
+if HAS_PYMONGO3:
+    from pymongo import MongoClient
+else:
+    from pymongo import Connection as MongoClient  # pylint: disable=E0611
 
 def compute_signature(index):
     signature = index["ns"]
@@ -18,7 +27,7 @@ def compute_signature(index):
     return signature
 
 def get_collection_stats(database, collection):
-    print "Checking DB: %s" % collection.full_name
+    print("Checking DB: %s" % collection.full_name)
     return database.command("collstats", collection.name)
 
 # From http://www.5dollarwhitebox.org/drupal/node/84
@@ -70,28 +79,48 @@ def get_cli_options():
                       default="",
                       metavar="PASSWORD",
                       help="Admin password if authentication is enabled")
+    parser.add_option("--ssl-cert",
+                      dest="ssl_certfile",
+                      default=None,
+                      metavar="CERTIFICATE",
+                      help="SSL Certificate to use is SSL is enabled")
+    parser.add_option("--ssl-ca-certs",
+                      dest="ssl_ca_certs",
+                      default=None,
+                      metavar="CA",
+                      help="SSL Certificate of CA for certificate validation if SSL is enabled")
 
     (options, args) = parser.parse_args()
 
     return options
 
-def get_connection(host, port, username, password):
+def get_connection(host, port, username, password, ssl_certfile=None, ssl_ca_certs=None):
     userPass = ""
     if username and password:
         userPass = username + ":" + password + "@"
 
     mongoURI = "mongodb://" + userPass + host + ":" + str(port)
-    return Connection(host=mongoURI, read_preference=ReadPreference.SECONDARY)
 
-def main(options):
+    conn_kwargs = dict(host=mongoURI, read_preference=ReadPreference.SECONDARY)
+
+    if HAS_PYMONGO3:
+        conn_kwargs.update(dict(ssl_certfile=ssl_certfile, ssl_ca_certs=ssl_ca_certs))
+
+    return MongoClient(**conn_kwargs)
+
+def main(options=None):
+    if options is None:
+        options = get_cli_options()
+
     summary_stats = {
         "count" : 0,
         "size" : 0,
         "indexSize" : 0
     }
     all_stats = []
 
-    connection = get_connection(options.host, options.port, options.user, options.password)
+    connection = get_connection(options.host, options.port, options.user, options.password,
+                                options.ssl_certfile, options.ssl_ca_certs)
 
     all_db_stats = {}
 
@@ -141,11 +170,11 @@ def main(options):
                 x.add_row(row)
 
 
-    print "Index Overview"
-    print x.get_string(sortby="Collection")
+    print("Index Overview")
+    print(x.get_string(sortby="Collection"))
 
     print
-    print "Top 5 Largest Indexes"
+    print("Top 5 Largest Indexes")
     x = PrettyTable(["Collection", "Index","% Size", "Index Size"])
     x.align["Collection"] = "l"
     x.align["Index"] = "l"
@@ -156,19 +185,19 @@ def main(options):
     top_five_indexes = sorted(index_size_mapping.keys(), reverse=True)[0:5]
     for size in top_five_indexes:
         x.add_row(index_size_mapping.get(size))
-    print x
+    print(x)
     print
 
-    print "Total Documents:", summary_stats["count"]
-    print "Total Data Size:", convert_bytes(summary_stats["size"])
-    print "Total Index Size:", convert_bytes(summary_stats["indexSize"])
+    print("Total Documents: %s" % summary_stats["count"])
+    print("Total Data Size: %s" % convert_bytes(summary_stats["size"]))
+    print("Total Index Size: %s" % convert_bytes(summary_stats["indexSize"]))
 
     # this is only meaningful if we're running the script on localhost
-    if options.host == "localhost":
+    if options.host == "localhost" or options.host == getfqdn():
         ram_headroom = psutil.phymem_usage()[0] - summary_stats["indexSize"]
-        print "RAM Headroom:", convert_bytes(ram_headroom)
-        print "RAM Used: %s (%s%%)" % (convert_bytes(psutil.phymem_usage()[1]), psutil.phymem_usage()[3])
-        print "Available RAM Headroom:", convert_bytes((100 - psutil.phymem_usage()[3]) / 100 * ram_headroom)
+        print("RAM Headroom: %s" % convert_bytes(ram_headroom))
+        print("RAM Used: %s (%s%%)" % (convert_bytes(psutil.phymem_usage()[1]), psutil.phymem_usage()[3]))
+        print("Available RAM Headroom: %s" % convert_bytes((100 - psutil.phymem_usage()[3]) / 100 * ram_headroom))
 
 if __name__ == "__main__":
     options = get_cli_options()

mongodbtools/redundant_indexes.py

@@ -6,10 +6,17 @@
 with just fields {field1:1}, the latter index is not needed since the first index already
 indexes the necessary fields.
 """
-from pymongo import Connection
 from pymongo import ReadPreference
 from optparse import OptionParser
+from distutils.version import StrictVersion
+import pymongo
 
+HAS_PYMONGO3 = bool(StrictVersion(pymongo.version) >= StrictVersion('3.0'))
+
+if HAS_PYMONGO3:
+    from pymongo import MongoClient
+else:
+    from pymongo import Connection as MongoClient  # pylint: disable=E0611
 
 def get_cli_options():
     parser = OptionParser(usage="usage: python %prog [options]",
@@ -40,21 +47,41 @@ def get_cli_options():
                       default="",
                       metavar="PASSWORD",
                       help="Admin password if authentication is enabled")
+    parser.add_option("--ssl-cert",
+                      dest="ssl_certfile",
+                      default=None,
+                      metavar="CERTIFICATE",
+                      help="SSL Certificate to use is SSL is enabled (only with pymongo >= 3)")
+    parser.add_option("--ssl-ca-certs",
+                      dest="ssl_ca_certs",
+                      default=None,
+                      metavar="CA",
+                      help="SSL Certificate of CA for certificate validation if SSL is enabled (only with pymongo >= 3)")
 
     (options, args) = parser.parse_args()
 
     return options
 
-def get_connection(host, port, username, password):
+def get_connection(host, port, username, password, ssl_certfile=None, ssl_ca_certs=None):
     userPass = ""
     if username and password:
         userPass = username + ":" + password + "@"
 
     mongoURI = "mongodb://" + userPass + host + ":" + str(port)
-    return Connection(host=mongoURI, read_preference=ReadPreference.SECONDARY)
 
-def main(options):
-    connection = get_connection(options.host, options.port, options.user, options.password)
+    conn_kwargs = dict(host=mongoURI, read_preference=ReadPreference.SECONDARY)
+
+    if HAS_PYMONGO3:
+        conn_kwargs.update(dict(ssl_certfile=ssl_certfile, ssl_ca_certs=ssl_ca_certs))
+
+    return MongoClient(**conn_kwargs)
+
+def main(options=None):
+    if options is None:
+        options = get_cli_options()
+
+    connection = get_connection(options.host, options.port, options.user, options.password,
+                                options.ssl_certfile, options.ssl_ca_certs)
 
     def compute_signature(index):
         signature = index["ns"]
@@ -66,7 +93,7 @@ def compute_signature(index):
         return signature
 
     def report_redundant_indexes(current_db):
-        print "Checking DB: %s" % current_db.name
+        print("Checking DB: %s" % current_db.name)
         indexes = current_db.system.indexes.find()
         index_map = {}
         for index in indexes:
@@ -78,11 +105,11 @@ def report_redundant_indexes(current_db):
                 if signature == other_sig:
                     continue
                 if other_sig.startswith(signature):
-                    print "Index %s[%s] may be redundant with %s[%s]" % (
+                    print ("Index %s[%s] may be redundant with %s[%s]" % (
                         index_map[signature]["ns"],
                         index_map[signature]["name"],
                         index_map[other_sig]["ns"],
-                        index_map[other_sig]["name"])
+                        index_map[other_sig]["name"]))
 
     databases= []
     if options.database:

requirements.txt

@@ -1,4 +1,4 @@
-pymongo==2.1
+pymongo>=2.1
 PrettyTable==0.7.1
 psutil==0.3.0
 mongoengine==0.5.0