simplify creat_table interface

Browse Source
This commit is contained in:
Yi Xie 2013-09-13 09:42:35 -04:00
parent bfd75360be
commit 95cc5dd411
3 changed files with 19 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
dataset/__init__.py
View File

@ -7,8 +7,9 @@ warnings.filterwarnings(
from dataset.persistence.database import Database
from dataset.persistence.table import Table
from dataset.freeze.app import freeze
from sqlalchemy import Integer, String
__all__ = ['Database', 'Table', 'freeze', 'connect']
__all__ = ['Database', 'Table', 'freeze', 'connect', 'Integer', 'String']
def connect(url=None, schema=None, reflectMetadata=True):

24
dataset/persistence/database.py
View File

@ -11,7 +11,7 @@ from sqlalchemy import Integer, String
from dataset.persistence.table import Table
from dataset.persistence.util import ResultIter
from dataset.util import DatasetException
log = logging.getLogger(__name__)
@ -97,14 +97,14 @@ class Database(object):
return list(set(self.metadata.tables.keys() +
self._tables.keys()))
def create_table(self, table_name, primary_id=None, is_integer=False):
def create_table(self, table_name, primary_id='id', primary_type=Integer):
"""
Creates a new table. The new table will automatically have an `id` column
unless specified via optional parameter primary_id, which will be used
as the primary key of the table. Automatic id is set to be an
auto-incrementing integer, while the type of custom primary_id can be a
string (i.e. VARCHAR(255)) by default or an integer via integer flag.
The caller will be responsible for the uniqueness of manually specified primary_id.
String (i.e. VARCHAR(255)) or an Integer as specified with primary_type flag.
The caller will be responsible for the uniqueness of manual primary_id.
This custom id feature is only available via direct create_table call.
@ -112,18 +112,22 @@ class Database(object):
::
table = db.create_table('population')
table2 = db.create_table('population2', primary_id='age', primary_type=Integer)
table3 = db.create_table('population3', primary_id='race', primary_type=String)
"""
self._acquire()
try:
log.debug("Creating table: %s on %r" % (table_name, self.engine))
table = SQLATable(table_name, self.metadata)
if not primary_id:
col = Column('id', Integer, primary_key=True)
else:
if is_integer:
col = Column(primary_id, Integer, primary_key=True, autoincrement=False)
else:
if primary_type is Integer:
auto_flag = False
if primary_id is 'id':
auto_flag = True
col = Column(primary_id, Integer, primary_key=True, autoincrement=auto_flag)
elif primary_type is String:
col = Column(primary_id, String(255), primary_key=True)
else:
raise DatasetException("The primary_type has to be either int or str.")
table.append_column(col)
table.create(self.engine)

6
test/test_persistence.py
View File

@ -2,7 +2,7 @@ import os
import unittest
from datetime import datetime
from dataset import connect
from dataset import connect, Integer, String
from dataset.util import DatasetException
from sample_data import TEST_DATA
from sqlalchemy.exc import IntegrityError
@ -35,7 +35,7 @@ class DatabaseTestCase(unittest.TestCase):
def test_create_table_custom_id1(self):
pid = "string_id"
table = self.db.create_table("foo2", primary_id = pid)
table = self.db.create_table("foo2", primary_id = pid, primary_type=String)
assert table.table.exists()
assert len(table.table.columns) == 1, table.table.columns
assert pid in table.table.c, table.table.c
@ -46,7 +46,7 @@ class DatabaseTestCase(unittest.TestCase):
def test_create_table_custom_id2(self):
pid = "int_id"
table = self.db.create_table("foo3", primary_id = pid, is_integer=True)
table = self.db.create_table("foo3", primary_id = pid, primary_type=Integer)
assert table.table.exists()
assert len(table.table.columns) == 1, table.table.columns
assert pid in table.table.c, table.table.c