Insert, update, expression and replace helpers

Summary:
Helpers for generating insert, update, and replace statements. Does perform column or table name escaping.

Test Plan: pytest

Reviewers: carl, neil

Differential Revision: http://grizzly.memsql.com/D5050

memsql/common/query_builder.py

@@ -0,0 +1,59 @@
+def simple_expression(joiner=', ', **fields):
+    """ Build a simple expression ready to be added onto another query.
+
+    >>> simple_expression(joiner=' AND ', name='bob', role='admin')
+    "`name`=%(_QB_name)s AND `name`=%(_QB_role)s", { '_QB_name': 'bob', '_QB_role': 'admin' }
+    """
+    expression, params = [], {}
+
+    for field_name, value in sorted(fields.iteritems(), key=lambda (n, v): n):
+        key = '_QB_%s' % field_name
+        expression.append('`%s`=%%(%s)s' % (field_name, key))
+        params[key] = value
+
+    return joiner.join(expression), params
+
+def update(table_name, **fields):
+    """ Build a update query.
+
+    >>> update('foo_table', a=5, b=2)
+    "UPDATE `foo_table` SET `a`=%(_QB_a)s, `b`=%(_QB_b)s", { '_QB_a': 5, '_QB_b': 2 }
+    """
+    prefix = "UPDATE `%s` SET " % table_name
+    sets, params = simple_expression(', ', **fields)
+    return prefix + sets, params
+
+def multi_insert(table_name, *rows):
+    """ Build a multi-insert query.
+        Each row in rows should be a dict of { column_name: column_value }
+
+    >>> multi_insert('foo_table', { 'a': 5, 'b': 2 }, { 'a': 5, 'b': 2 })
+    "INSERT INTO `foo_table` (`a`, `b`) VALUES (%(_QB_ROW_0)s), (%(_QB_ROW_1)s)"
+    """
+    return __multi_insert(table_name, rows)
+
+def multi_replace(table_name, *rows):
+    """ Build a multi-replace query.
+        Each row in rows should be a dict of { column_name: column_value }
+
+    >>> multi_replace('foo_table', { 'a': 5, 'b': 2 }, { 'a': 5, 'b': 2 })
+    "REPLACE INTO `foo_table` (`a`, `b`) VALUES (%(_QB_ROW_0)s), (%(_QB_ROW_1)s)"
+    """
+    return __multi_insert(table_name, rows, replace=True)
+
+def __multi_insert(table_name, rows, replace=False):
+    cols = sorted(rows[0].keys())
+    prefix = '%s INTO `%s` (%s) VALUES ' % (
+        'REPLACE' if replace else 'INSERT',
+        table_name,
+        ', '.join(['`%s`' % col for col in cols])
+    )
+    sql, params = [], {}
+
+    for i, row in enumerate(rows):
+        key = '_QB_ROW_%d' % i
+        params[key] = [ v for c, v in sorted(row.iteritems(), key=lambda (c, v): cols.index(c)) ]
+        sql.append('(%%(%s)s)' % key)
+
+    return prefix + ', '.join(sql), params
+

memsql/common/test/test_query_builder.py

@@ -0,0 +1,29 @@
+from memsql.common import query_builder, database
+
+def test_simple_expression():
+    x = { 'a': 1, 'b': '2', 'c': 1223.4 }
+    sql, params = query_builder.simple_expression(', ', **x)
+    assert sql == '`a`=%(_QB_a)s, `b`=%(_QB_b)s, `c`=%(_QB_c)s'
+    assert params == { '_QB_a': 1, '_QB_b': '2', '_QB_c': 1223.4 }
+    assert database.escape_query(sql, params) == r"`a`=1, `b`='2', `c`=1223.4"
+
+def test_update():
+    x = { 'a': 1, 'b': '2', 'c': 1223.4 }
+    sql, params = query_builder.update('foo', **x)
+    assert sql == 'UPDATE `foo` SET `a`=%(_QB_a)s, `b`=%(_QB_b)s, `c`=%(_QB_c)s'
+    assert params == { '_QB_a': 1, '_QB_b': '2', '_QB_c': 1223.4 }
+    assert database.escape_query(sql, params) == r"UPDATE `foo` SET `a`=1, `b`='2', `c`=1223.4"
+
+def test_multi_insert():
+    rows = [{ 'a': 1, 'b': '2', 'c': 1223.4 }, { 'a': 2, 'b': '5', 'c': 1 }]
+    sql, params = query_builder.multi_insert('foo', *rows)
+    assert sql == 'INSERT INTO `foo` (`a`, `b`, `c`) VALUES (%(_QB_ROW_0)s), (%(_QB_ROW_1)s)'
+    assert params == { '_QB_ROW_0': [1, '2', 1223.4], '_QB_ROW_1': [2, '5', 1] }
+    assert database.escape_query(sql, params) == r"INSERT INTO `foo` (`a`, `b`, `c`) VALUES (1,'2',1223.4), (2,'5',1)"
+
+def test_replace():
+    rows = [{ 'a': 1, 'b': '2', 'c': 1223.4 }, { 'a': 2, 'b': '5', 'c': 1 }]
+    sql, params = query_builder.multi_replace('foo', *rows)
+    assert sql == 'REPLACE INTO `foo` (`a`, `b`, `c`) VALUES (%(_QB_ROW_0)s), (%(_QB_ROW_1)s)'
+    assert params == { '_QB_ROW_0': [1, '2', 1223.4], '_QB_ROW_1': [2, '5', 1] }
+    assert database.escape_query(sql, params) == r"REPLACE INTO `foo` (`a`, `b`, `c`) VALUES (1,'2',1223.4), (2,'5',1)"