mediawiki_dump_tools/Mediawiki-Utilities/mw/database/collections/users.py

   1 import logging
   2 import time
   3
   4 from ...types import Timestamp
   5 from ...util import none_or
   6 from .collection import Collection
   7
   8 logger = logging.getLogger("mw.database.collections.users")
   9
  10
  11 class Users(Collection):
  12     CREATION_ACTIONS = {'newusers', 'create', 'create2', 'autocreate',
  13                         'byemail'}
  14
  15     def get(self, user_id=None, user_name=None):
  16         """
  17         Gets a single user row from the database.  Raises a :class:`KeyError`
  18         if a user cannot be found.
  19
  20         :Parameters:
  21             user_id : int
  22                 User ID
  23             user_name : str
  24                 User's name
  25
  26         :Returns:
  27             A user row.
  28         """
  29         user_id = none_or(user_id, int)
  30         user_name = none_or(user_name, str)
  31
  32         query = """
  33             SELECT user.*
  34             FROM user
  35         """
  36         values = []
  37
  38         if user_id is not None:
  39             query += """
  40                 WHERE user_id = %s
  41             """
  42             values.append(user_id)
  43
  44         elif user_name is not None:
  45             query += """
  46                 WHERE user_name = %s
  47             """
  48             values.append(user_name)
  49
  50         else:
  51             raise TypeError("Must specify a user identifier.")
  52
  53         cursor = self.db.shared_connection.cursor()
  54         cursor.execute(
  55             query,
  56             values
  57         )
  58
  59         for row in cursor:
  60             return row
  61
  62         raise KeyError(user_id if user_id is not None else user_name)
  63
  64     def query(self, registered_before=None, registered_after=None,
  65               before_id=None, after_id=None, limit=None,
  66               direction=None, self_created_only=False):
  67         """
  68         Queries users based on various filtering parameters.
  69
  70         :Parameters:
  71             registered_before : :class:`mw.Timestamp`
  72                 A timestamp to search before (inclusive)
  73             registered_after : :class:`mw.Timestamp`
  74                 A timestamp to search after (inclusive)
  75             before_id : int
  76                 A user_id to search before (inclusive)
  77             after_id : int
  78                 A user_ud to search after (inclusive)
  79             direction : str
  80                 "newer" or "older"
  81             limit : int
  82                 Limit the results to at most this number
  83             self_creations_only : bool
  84                 limit results to self_created user accounts
  85
  86         :Returns:
  87             an iterator over ``user`` table rows
  88         """
  89         start_time = time.time()
  90
  91         registered_before = none_or(registered_before, Timestamp)
  92         registered_after = none_or(registered_after, Timestamp)
  93         before_id = none_or(before_id, str)
  94         after_id = none_or(after_id, str)
  95         direction = none_or(direction, levels=self.DIRECTIONS)
  96         limit = none_or(limit, int)
  97         self_created_only = bool(self_created_only)
  98
  99         query = """
 100             SELECT user.*
 101             FROM user
 102         """
 103         values = []
 104
 105         if self_created_only:
 106             query += """
 107                 INNER JOIN logging ON
 108                     log_user = user_id
 109                     log_type = "newusers" AND
 110                     log_action = "create"
 111             """
 112
 113         query += "WHERE 1 "
 114
 115         if registered_before is not None:
 116             query += "AND user_registration <= %s "
 117             values.append(registered_before.short_format())
 118         if registered_after is not None:
 119             query += "AND user_registration >= %s "
 120             values.append(registered_after.short_format())
 121         if before_id is not None:
 122             query += "AND user_id <= %s "
 123             values.append(before_id)
 124         if after_id is not None:
 125             query += "AND user_id >= %s "
 126             values.append(after_id)
 127
 128         query += "GROUP BY user_id "  # In case of duplicate log events
 129
 130         if direction is not None:
 131             if registered_before is not None or registered_after is not None:
 132                 if direction == "newer":
 133                     query += "ORDER BY user_registration ASC "
 134                 else:
 135                     query += "ORDER BY user_registration DESC "
 136             else:
 137                 if direction == "newer":
 138                     query += "ORDER BY user_id ASC "
 139                 else:
 140                     query += "ORDER BY user_id DESC "
 141
 142         if limit is not None:
 143             query += "LIMIT %s "
 144             values.append(limit)
 145
 146         cursor = self.db.shared_connection.cursor()
 147         cursor.execute(query, values)
 148
 149         count = 0
 150         for row in cursor:
 151             yield row
 152             count += 1
 153
 154         logger.debug("%s users queried in %s seconds" % (count, time.time() - start_time))