Dd1UwFsy

1from __future__ import print_function
2
3class DbNotFound(Exception):
4    pass
5
6def copy(src_conn, src_dbname, dst_conn, dst_dbname=None, drop_table=True,
7    drop_db=True, silent=False, tables_to_copy=None):
8    """ (<lurker connection>, str, <lurker connection>, dst_dbname=str, drop_table=bool
9        drop_db=bool, silent=bool, tables_to_copy=list) -> bool
10
11        Copies database (src_dbname) from source connection (src_conn) to destination
12        database (dst_dbname) at destination connection (dst_conn)
13
14        If dst_dbname is not given, will be same as src_dbname.
15    """
16
17    def values2Str(values_dict):
18        values_tuple_string = u"("
19        for key, value in values_dict.iteritems():
20            if value == None:
21                values_tuple_string += u","
22            elif type(value) in (unicode, datetime.date):
23                values_tuple_string += u"'%s', " % value
24            elif type(value) == int or type(value) == long:
25                values_tuple_string += u"%s, " % value
26        values_tuple_string = values_tuple_string[:-2]
27        values_tuple_string += u")"
28        return unicode(values_tuple_string)
29
30    def keys2Str(values_dict):
31        keys_tuple_string = u"("
32        for key, value in values_dict.iteritems():
33            keys_tuple_string += "%s, " % key
34        keys_tuple_string = keys_tuple_string[:-2]
35        keys_tuple_string += u")"
36        return unicode(keys_tuple_string)
37
38    if not dst_dbname:
39        dst_dbname = src_dbname
40
41    # Get db names
42    db_names = [result['Database'] for result in \
43        src_conn.get_results("show databases")]
44
45    # If source db is not exists, raise
46    if not src_dbname in db_names:
47        raise DbNotFound("Database '%s' is not found at source connection" % (
48            src_dbname))
49
50    table_names = [result['Tables_in_%s' % src_dbname] for result in
51        src_conn.get_results("show full tables in %s where Table_Type ='BASE TABLE'" % src_dbname)]
52
53    # If there is spesific table names to copy, filter table_names
54    if tables_to_copy:
55        table_names = filter(lambda tn: tn in tables_to_copy, table_names)
56
57    dst_conn.execute("SET FOREIGN_KEY_CHECKS = 0")
58
59    if drop_db:
60        dst_conn.execute("DROP DATABASE IF EXISTS %s" % dst_dbname)
61
62    dst_conn.execute("CREATE DATABASE IF NOT EXISTS %s" % dst_dbname)
63
64    for table_name in table_names:
65
66        create_table_query = src_conn.get_results(
67            "show create table `%s`.`%s`" % (src_dbname, table_name)
68        )[0]['Create Table']
69
70        create_table_query = "%s %s.%s" % (
71            create_table_query[:12], dst_dbname, create_table_query[13:])
72
73        # Create table on destination
74        dst_conn.execute(
75            "drop table if exists `%s`.`%s`" % (dst_dbname, table_name))
76
77        dst_conn.execute(create_table_query)
78
79        # Get total rows of source table
80        num_of_rows = src_conn.get_row(
81            "select count(*) from `%s`.`%s`" % (src_dbname, table_name)
82        )['count(*)']
83
84        # Create loop for copying rows.
85        row_idx = 0
86
87        for values_dict in src_conn.get_results(
88            "select * from `%s`.`%s`" % (src_dbname, table_name)):
89
90            # Create inserting query
91            query = u"INSERT INTO `%s`.`%s` %s VALUES %s" % (
92                dst_dbname, table_name, keys2Str(values_dict),
93                values2Str(values_dict))
94
95            row_idx += 1
96
97            if not silent:
98                # In every 1000 row, print status
99                if row_idx % 1000 == 0 or row_idx == num_of_rows:
100                    print("[%s] Copying table: %s [%s/%s rows completed]" % (\
101                        datetime.now(), table_name, row_idx, num_of_rows))
102
103            # Execute query
104            dst_conn.execute(query)
105
106        if num_of_rows and not silent:
107            print()