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.database import Database
from dataset.persistence.table import Table from dataset.persistence.table import Table
from dataset.freeze.app import freeze 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): 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.table import Table
from dataset.persistence.util import ResultIter from dataset.persistence.util import ResultIter
from dataset.util import DatasetException
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@ -97,14 +97,14 @@ class Database(object):
return list(set(self.metadata.tables.keys() + return list(set(self.metadata.tables.keys() +
self._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 Creates a new table. The new table will automatically have an `id` column
unless specified via optional parameter primary_id, which will be used 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 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 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. String (i.e. VARCHAR(255)) or an Integer as specified with primary_type flag.
The caller will be responsible for the uniqueness of manually specified primary_id. The caller will be responsible for the uniqueness of manual primary_id.
This custom id feature is only available via direct create_table call. This custom id feature is only available via direct create_table call.
@ -112,18 +112,22 @@ class Database(object):
:: ::
table = db.create_table('population') 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() self._acquire()
try: try:
log.debug("Creating table: %s on %r" % (table_name, self.engine)) log.debug("Creating table: %s on %r" % (table_name, self.engine))
table = SQLATable(table_name, self.metadata) table = SQLATable(table_name, self.metadata)
if not primary_id: if primary_type is Integer:
col = Column('id', Integer, primary_key=True) auto_flag = False
else: if primary_id is 'id':
if is_integer: auto_flag = True
col = Column(primary_id, Integer, primary_key=True, autoincrement=False) col = Column(primary_id, Integer, primary_key=True, autoincrement=auto_flag)
else: elif primary_type is String:
col = Column(primary_id, String(255), primary_key=True) 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.append_column(col)
table.create(self.engine) table.create(self.engine)

6
test/test_persistence.py
View File

@ -2,7 +2,7 @@ import os
import unittest import unittest
from datetime import datetime from datetime import datetime
from dataset import connect from dataset import connect, Integer, String
from dataset.util import DatasetException from dataset.util import DatasetException
from sample_data import TEST_DATA from sample_data import TEST_DATA
from sqlalchemy.exc import IntegrityError from sqlalchemy.exc import IntegrityError
@ -35,7 +35,7 @@ class DatabaseTestCase(unittest.TestCase):
def test_create_table_custom_id1(self): def test_create_table_custom_id1(self):
pid = "string_id" 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 table.table.exists()
assert len(table.table.columns) == 1, table.table.columns assert len(table.table.columns) == 1, table.table.columns
assert pid in table.table.c, table.table.c assert pid in table.table.c, table.table.c
@ -46,7 +46,7 @@ class DatabaseTestCase(unittest.TestCase):
def test_create_table_custom_id2(self): def test_create_table_custom_id2(self):
pid = "int_id" 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 table.table.exists()
assert len(table.table.columns) == 1, table.table.columns assert len(table.table.columns) == 1, table.table.columns
assert pid in table.table.c, table.table.c assert pid in table.table.c, table.table.c