add function: make_sharding

Author: wenbobuaa

mysqlutil/README.md

@@ -17,6 +17,7 @@
   - [mysqlutil.make_insert_sql](#mysqlutilmake_insert_sql)
   - [mysqlutil.make_range_mysqldump_cmd](#mysqlutilmake_range_mysqldump_cmd)
   - [mysqlutil.make_select_sql](#mysqlutilmake_select_sql)
+  - [mysqlutil.make_sharding](#mysqlutilmake_sharding)
   - [mysqlutil.make_sql_range_conditions](#mysqlutilmake_sql_range_conditions)
   - [mysqlutil.make_update_sql](#mysqlutilmake_update_sql)
   - [mysqlutil.scan_index](#mysqlutilscan_index)
@@ -423,6 +424,65 @@ make_select_sql('errlog', ['_id', 'key'], ('key', 'val'), ('a', 'b'),
 a string which is a sql select statement.
 
 
+## mysqlutil.make_sharding
+
+**syntax**:
+`mysqlutil.make_sharding(conf)`
+
+Scan a database table and generate sharding info according configurations in `conf`.
+Return sharding result as a dictionary like:
+```
+{
+    "shard": [(), (), ...],
+    "number": [number, number, ...],
+    "total": number,
+}
+```
+
+**argument**:
+
+-   `db`: which database to sharding. A string.
+-   `table`: which table to sharding. A string.
+-   `conn`: database connect info:
+
+    ```
+    {
+        'host': '127.0.0.1',
+        'port': 3306,
+        'user': 'mysql',
+        'passwd': 'password',
+    }
+    ```
+
+-   `shard_fields`: are index fields to sharding by, a list or tuple of strings.
+-   `start`: is the start condition to scan table, a list or tuple of strings.
+-   `number_per_shard`: specifies the number of rows a shard contains, an integer.
+-   `tolerance_of_shard`: the tolerance of one shard's capacity, an integer.
+-   `shard_maker`: a function which accepts one list of strings argument and return a value as
+    a "shard" in the result. For example:
+
+    ```
+    def shard_maker(shard):
+        new_shard = shard + ['']*3
+        return new_shard[:3]
+    ```
+    By default, it is `list`.
+
+**return**:
+a dictionary of sharding result, like:
+```
+{
+    "shard": [['10000', 'a', '1'], ['11000', 'b', '3'], ['12000', 'd', '9']],
+    "number": [100, 104, 80],
+    "total": 284,
+}
+```
+
+-   `shard`: sharding info, a list of first row of every shard.
+-   `number`: numbers of rows of every shard.
+-   `total`: number of rows of all shards.
+
+
 ## mysqlutil.make_sql_range_conditions
 
 **syntax**:

mysqlutil/__init__.py

@@ -11,6 +11,7 @@
     make_insert_sql,
     make_range_mysqldump_cmd,
     make_select_sql,
+    make_sharding,
     make_sql_range_conditions,
     make_update_sql,
     scan_index,
@@ -35,6 +36,7 @@
     "make_insert_sql",
     "make_range_mysqldump_cmd",
     "make_select_sql",
+    "make_sharding",
     "make_sql_range_conditions",
     "make_update_sql",
     "privileges",

mysqlutil/mysqlutil.py

@@ -8,6 +8,7 @@
 
 from pykit import dictutil
 from pykit import mysqlconnpool
+from pykit import strutil
 
 
 class ConnectionTypeError(Exception):
@@ -375,3 +376,32 @@ def make_where_clause(index, index_values, operator='='):
         where_clause = ''
 
     return where_clause
+
+
+def make_sharding(db, table, conn, shard_fields, start, number_per_shard, tolerance_of_shard,
+        shard_maker=list):
+
+    scan_args = [(db, table), shard_fields, shard_fields, start]
+    scan_kwargs = {"left_open": False, "use_dict": False, "retry": 3}
+
+    connpool = mysqlconnpool.make(conn)
+
+    record_iter = scan_index(connpool, *scan_args, **scan_kwargs)
+
+    shards = strutil.sharding(
+        record_iter, number_per_shard, accuracy=tolerance_of_shard, joiner=list)
+
+    _, count = shards[0]
+    result = {
+        "shard": [shard_maker(start)],
+        "number": [count],
+        "total": count,
+    }
+
+    for shard, count in shards[1:]:
+
+        result['shard'].append(shard_maker(shard))
+        result['number'].append(count)
+        result['total'] += count
+
+    return result

mysqlutil/test/test_mysqlutil.py

@@ -119,6 +119,114 @@ def test_scan_index(self):
             except error as e:
                 self.assertEqual(type(e), error)
 
+    def test_make_sharding(self):
+
+        db = mysql_test_db
+        table = mysql_test_table
+        conn = {
+            'host': base.mysql_test_ip,
+            'port': base.mysql_test_port,
+            'user': mysql_test_user,
+            'passwd': base.mysql_test_password,
+        }
+
+        def shard_maker(shard):
+
+            new_shard = [str(x) for x in shard]
+            new_shard += ['', '', '']
+
+            return tuple(new_shard[:3])
+
+        cases = (
+            (
+                {
+                    "shard_fields": ('service', 'ip', '_id'),
+                    "start": ['common0', '', ''],
+                    "number_per_shard": 10,
+                    "tolerance_of_shard": 1,
+                    "shard_maker": tuple,
+                },
+                {
+                    'total': 32,
+                    'number': [10, 10, 10, 2],
+                    'shard': [('common0', '', ''), ('common0', '127.0.0.3', 27L),
+                              ('common2', '127.0.0.1'), ('common4', '127.0.0.1', 7L)],
+                },
+            ),
+
+            (
+                {
+                    "shard_fields": ('service', 'ip', '_id'),
+                    "start": ['common0', '127.0.0.3', '27'],
+                    "number_per_shard": 10,
+                    "tolerance_of_shard": 1,
+                },
+                {
+                    'total': 22,
+                    'number': [10, 10, 2],
+                    'shard': [['common0', '127.0.0.3', '27'],
+                              ['common2', '127.0.0.1'], ['common4', '127.0.0.1', 7L]],
+                },
+            ),
+
+            (
+                {
+                    "shard_fields": ('service', 'ip', '_id'),
+                    "start": ['common0', '', ''],
+                    "number_per_shard": 10,
+                    "tolerance_of_shard": 1,
+                    "shard_maker": shard_maker,
+                },
+                {
+                    'total': 32,
+                    'number': [10, 10, 10, 2],
+                    'shard': [('common0', '', ''), ('common0', '127.0.0.3', '27'),
+                              ('common2', '127.0.0.1', ''), ('common4', '127.0.0.1', '7')],
+                },
+            ),
+
+            (
+                {
+                    "shard_fields": ('service', 'ip', '_id'),
+                    "start": ['common0', '', ''],
+                    "number_per_shard": 15,
+                    "tolerance_of_shard": 2,
+                    "shard_maker": shard_maker,
+                },
+                {
+                    'total': 32,
+                    'number': [15, 15, 2],
+                    'shard': [('common0', '', ''), ('common1', '127.0.0.1', '31'), ('common4', '', '')],
+                },
+            ),
+
+            (
+                {
+                    "shard_fields": ('time', '_id'),
+                    "start": ['201706060600', '1'],
+                    "number_per_shard": 10,
+                    "tolerance_of_shard": 1,
+                },
+                {
+                    'total': 32,
+                    'number': [10, 10, 10, 2],
+                    'shard': [['201706060600', '1'], [201706060610L,], [201706060620L,],
+                              [201706060630L,]],
+                },
+            ),
+        )
+
+        for kwargs, expected in cases:
+
+            kwargs['db'] = db
+            kwargs['table'] = table
+            kwargs['conn'] = conn
+
+            dd('expected: ', expected)
+            result = mysqlutil.make_sharding(**kwargs)
+            dd('result  : ', result)
+            self.assertEqual(result, expected)
+
 
 class TestMysqlutil(unittest.TestCase):