mediawiki_dump_tools/Mediawiki-Utilities/mw/database/db.py

   1 import getpass
   2 import logging
   3 import os
   4
   5 import pymysql
   6 import pymysql.cursors
   7
   8 from .collections import AllRevisions, Archives, Pages, Revisions, Users
   9
  10 logger = logging.getLogger("mw.database.db")
  11
  12
  13 class DB:
  14     """
  15     Represents a connection to a MySQL database.
  16
  17     :Parameters:
  18         connection = :class:`oursql.Connection`
  19             A connection to a MediaWiki database
  20     """
  21
  22     def __init__(self, connection):
  23         self.shared_connection = connection
  24         self.shared_connection.cursorclass = pymysql.cursors.DictCursor
  25
  26         self.revisions = Revisions(self)
  27         """
  28         An instance of :class:`mw.database.Revisions`.
  29         """
  30
  31         self.archives = Archives(self)
  32         """
  33         An instance of :class:`mw.database.Archives`.
  34         """
  35
  36         self.all_revisions = AllRevisions(self)
  37         """
  38         An instance of :class:`mw.database.AllRevisions`.
  39         """
  40
  41         self.pages = Pages(self)
  42         """
  43         An instance of :class:`mw.database.Pages`.
  44         """
  45
  46         self.users = Users(self)
  47         """
  48         An instance of :class:`mw.database.Users`.
  49         """
  50
  51     def __repr__(self):
  52         return "%s(%s)" % (
  53             self.__class__.__name__,
  54             ", ".join(
  55                 [repr(arg) for arg in self.args] +
  56                 ["%s=%r" % (k, v) for k, v in self.kwargs.items()]
  57             )
  58         )
  59
  60     def __str__(self):
  61         return self.__repr__()
  62
  63     @classmethod
  64     def add_arguments(cls, parser, defaults=None):
  65         """
  66         Adds the arguments to an :class:`argparse.ArgumentParser` in order to
  67         create a database connection.
  68         """
  69         defaults = defaults if defaults is not None else defaults
  70
  71         default_host = defaults.get('host', "localhost")
  72         parser.add_argument(
  73             '--host', '-h',
  74             help="MySQL database host to connect to (defaults to {0})".format(default_host),
  75             default=default_host
  76         )
  77
  78         default_database = defaults.get('database', getpass.getuser())
  79         parser.add_argument(
  80             '--database', '-d',
  81             help="MySQL database name to connect to (defaults to  {0})".format(default_database),
  82             default=default_database
  83         )
  84
  85         default_defaults_file = defaults.get('defaults-file', os.path.expanduser("~/.my.cnf"))
  86         parser.add_argument(
  87             '--defaults-file',
  88             help="MySQL defaults file (defaults to {0})".format(default_defaults_file),
  89             default=default_defaults_file
  90         )
  91
  92         default_user = defaults.get('user', getpass.getuser())
  93         parser.add_argument(
  94             '--user', '-u',
  95             help="MySQL user (defaults to %s)".format(default_user),
  96             default=default_user
  97         )
  98         return parser
  99
 100     @classmethod
 101     def from_arguments(cls, args):
 102         """
 103         Constructs a :class:`~mw.database.DB`.
 104         Consumes :class:`argparse.ArgumentParser` arguments given by
 105         :meth:`add_arguments` in order to create a :class:`DB`.
 106
 107         :Parameters:
 108             args : :class:`argparse.Namespace`
 109                 A collection of argument values returned by :class:`argparse.ArgumentParser`'s :meth:`parse_args()`
 110         """
 111         connection = pymysql.connect(
 112             args.host,
 113             args.user,
 114             database=args.database,
 115             read_default_file=args.defaults_file
 116         )
 117         return cls(connection)
 118
 119     @classmethod
 120     def from_params(cls, *args, **kwargs):
 121         """
 122         Constructs a :class:`~mw.database.DB`.  Passes `*args` and `**kwargs`
 123         to :meth:`oursql.connect` and configures the connection.
 124
 125         :Parameters:
 126             args : :class:`argparse.Namespace`
 127                 A collection of argument values returned by :class:`argparse.ArgumentParser`'s :meth:`parse_args()`
 128         """
 129         kwargs['cursorclass'] = pymysql.cursors.DictCursor
 130         if kwargs['db']:
 131             kwargs['database'] = kwargs['db']
 132             del kwargs['db']
 133         connection = pymysql.connect(*args, **kwargs)
 134         return cls(connection)