Sf4yXxtU

1"""
2Functions for generating SQL statements related to user-data 
3database tables and rows.
4"""
5
6_default_columns = [
7    'id SERIAL',
8    'CONSTRAINT primary_key PRIMARY KEY (id)',
9]
10
11_sql_templates = {
12    'create_table': 'CREATE TABLE %(name_fragment)s (%(columns_fragment)s);',
13    'alter_table': 'ALTER TABLE %(name_fragment)s %(columns_fragment)s;',
14    'drop_table': 'DROP TABLE %(name_fragment)s;',
15}
16
17
18def create_table(name, columns=None):
19    """
20    Returns a ``CREATE TABLE`` statement for the given table name and 
21    list of column definitions, as created by the ``column`` function. New 
22    tables will always include the columns and constraints included 
23    in ``_default_columns``, most importantly a primary key column.
24    
25    TODO: Ensure column names don't collide with default columns.
26
27    >>> columns = [column('col_name1', 'integer'), column('col_name2', 'integer')]
28    >>> create_table('table_name', columns)
29    'CREATE TABLE table_name (id SERIAL, CONSTRAINT primary_key PRIMARY KEY (id), col_name1 INTEGER, col_name2 INTEGER);'
30    """
31    context = {
32        'name_fragment': name,
33        'columns_fragment': ', '.join(
34            _default_columns + columns if columns else _default_columns
35        ),
36    }
37
38    return _sql_templates['create_table'] % context
39
40
41def add_columns(name, columns):
42    """
43    Return an ALTER TABLE statement that will insert the given 
44    list of column definitions, as created by ``column`` 
45    into the table named ``name``.
46
47    >>> columns = [column('col_name1', 'integer'), column('col_name2', 'integer')]
48    >>> add_columns('table_name', columns)
49    'ALTER TABLE table_name ADD COLUMN col_name1 INTEGER, ADD COLUMN col_name2 INTEGER;'
50    """
51    context = {
52        'name_fragment': name,
53        'columns_fragment': ', '.join(
54            ['ADD COLUMN %s' % c for c in columns]
55        ),
56    }
57
58    return _sql_templates['alter_table'] % context
59
60
61def drop_columns(name, columns):
62    """
63    Return an ALTER TABLE statement that will drop the given 
64    columns, given as plain old column name from the table 
65    named ``name``.
66
67    >>> columns = ['col1', 'col2']
68    >>> drop_columns('table_name', columns)
69    'ALTER TABLE table_name DROP COLUMN IF EXISTS col1, DROP COLUMN IF EXISTS col2;'
70    """
71    context = {
72        'name_fragment': name,
73        'columns_fragment': ', '.join(
74            ['DROP COLUMN IF EXISTS %s' % c for c in columns]
75        ),
76    }
77
78    return _sql_templates['alter_table'] % context
79
80
81def drop_table(name):
82    """
83    Returns a DROP TABLE statement that drops the table with the 
84    given name.
85
86    >>> drop_table('table_name')
87    'DROP TABLE table_name;'
88    """
89    context = {
90        'name_fragment': name,
91    }
92    
93    return _sql_templates['drop_table'] % context
94    
95    
96def column(name, datatype, unique=False, required=False):
97    """
98    Returns a valid column definition for use with an ``ALTER TABLE`` 
99    or ``CREATE TABLE`` statement.
100    
101    TODO: Handle a default value.
102    
103    >>> column('col_name', 'integer', True, True)
104    'col_name INTEGER UNIQUE NOT NULL'
105    >>> column('col_name', 'integer', unique=True)
106    'col_name INTEGER UNIQUE'
107    >>> column('col_name', 'integer', required=True)
108    'col_name INTEGER NOT NULL'
109    >>> column('col_name', 'integer')
110    'col_name INTEGER'
111    """
112    return '%s %s%s%s' % (
113        name,
114        datatype.upper(),
115        ' UNIQUE' if unique else '',
116        ' NOT NULL' if required else '',
117    )
118
119
120if __name__ == '__main__':
121    import doctest
122    doctest.testmod()