From bfd75360bebfea6174b915737b5618b295a9b833 Mon Sep 17 00:00:00 2001 From: Yi Xie Date: Sun, 8 Sep 2013 11:35:43 -0400 Subject: [PATCH 1/3] create table with custom primary_id --- dataset/persistence/database.py | 23 ++++++++++++++++++----- test/test_persistence.py | 28 +++++++++++++++++++++++++++- 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/dataset/persistence/database.py b/dataset/persistence/database.py index 20ca061..20a1519 100644 --- a/dataset/persistence/database.py +++ b/dataset/persistence/database.py @@ -7,7 +7,7 @@ from migrate.versioning.util import construct_engine from sqlalchemy.pool import NullPool from sqlalchemy.schema import MetaData, Column, Index from sqlalchemy.schema import Table as SQLATable -from sqlalchemy import Integer +from sqlalchemy import Integer, String from dataset.persistence.table import Table from dataset.persistence.util import ResultIter @@ -97,10 +97,16 @@ class Database(object): return list(set(self.metadata.tables.keys() + self._tables.keys())) - def create_table(self, table_name): + def create_table(self, table_name, primary_id=None, is_integer=False): """ - Creates a new table. The new table will automatically have an `id` column, which is - set to be an auto-incrementing integer as the primary key of the table. + 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. + + This custom id feature is only available via direct create_table call. Returns a :py:class:`Table ` instance. :: @@ -111,7 +117,14 @@ class Database(object): try: log.debug("Creating table: %s on %r" % (table_name, self.engine)) table = SQLATable(table_name, self.metadata) - col = Column('id', Integer, primary_key=True) + 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: + col = Column(primary_id, String(255), primary_key=True) + table.append_column(col) table.create(self.engine) self._tables[table_name] = table diff --git a/test/test_persistence.py b/test/test_persistence.py index 893ceb3..80614b8 100644 --- a/test/test_persistence.py +++ b/test/test_persistence.py @@ -5,13 +5,14 @@ from datetime import datetime from dataset import connect from dataset.util import DatasetException from sample_data import TEST_DATA +from sqlalchemy.exc import IntegrityError class DatabaseTestCase(unittest.TestCase): def setUp(self): os.environ['DATABASE_URL'] = 'sqlite:///:memory:' - self.db = connect() + self.db = connect('sqlite:///:memory:') self.tbl = self.db['weather'] for row in TEST_DATA: self.tbl.insert(row) @@ -32,6 +33,31 @@ class DatabaseTestCase(unittest.TestCase): assert len(table.table.columns) == 1, table.table.columns assert 'id' in table.table.c, table.table.c + def test_create_table_custom_id1(self): + pid = "string_id" + table = self.db.create_table("foo2", primary_id = pid) + assert table.table.exists() + assert len(table.table.columns) == 1, table.table.columns + assert pid in table.table.c, table.table.c + + table.insert({ + 'string_id': 'foobar'}) + assert table.find_one(string_id = 'foobar')[0] == 'foobar' + + def test_create_table_custom_id2(self): + pid = "int_id" + table = self.db.create_table("foo3", primary_id = pid, is_integer=True) + assert table.table.exists() + assert len(table.table.columns) == 1, table.table.columns + assert pid in table.table.c, table.table.c + + table.insert({'int_id': 123}) + table.insert({'int_id': 124}) + assert table.find_one(int_id = 123)[0] == 123 + assert table.find_one(int_id = 124)[0] == 124 + with self.assertRaises(IntegrityError): + table.insert({'int_id': 123}) + def test_load_table(self): tbl = self.db.load_table('weather') assert tbl.table == self.tbl.table From 95cc5dd411f0ca690dd8c65e2a602664d547e22c Mon Sep 17 00:00:00 2001 From: Yi Xie Date: Fri, 13 Sep 2013 09:42:35 -0400 Subject: [PATCH 2/3] simplify creat_table interface --- dataset/__init__.py | 3 ++- dataset/persistence/database.py | 24 ++++++++++++++---------- test/test_persistence.py | 6 +++--- 3 files changed, 19 insertions(+), 14 deletions(-) diff --git a/dataset/__init__.py b/dataset/__init__.py index 4441aee..f6cf604 100644 --- a/dataset/__init__.py +++ b/dataset/__init__.py @@ -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): diff --git a/dataset/persistence/database.py b/dataset/persistence/database.py index 20a1519..1ae27b2 100644 --- a/dataset/persistence/database.py +++ b/dataset/persistence/database.py @@ -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) + 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: - if is_integer: - col = Column(primary_id, Integer, primary_key=True, autoincrement=False) - else: - col = Column(primary_id, String(255), primary_key=True) + raise DatasetException("The primary_type has to be either int or str.") table.append_column(col) table.create(self.engine) diff --git a/test/test_persistence.py b/test/test_persistence.py index 80614b8..063032c 100644 --- a/test/test_persistence.py +++ b/test/test_persistence.py @@ -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 From 257d767b6dccf5f4b7fb7800341e4d598a258cf9 Mon Sep 17 00:00:00 2001 From: Yi Xie Date: Sun, 15 Sep 2013 14:12:30 -0400 Subject: [PATCH 3/3] add options to create custom primary id and type with shorthand format --- dataset/__init__.py | 4 ++-- dataset/persistence/database.py | 36 ++++++++++++++++++++++----------- test/test_persistence.py | 31 +++++++++++++++++++++++++--- 3 files changed, 54 insertions(+), 17 deletions(-) diff --git a/dataset/__init__.py b/dataset/__init__.py index f6cf604..d9d4f96 100644 --- a/dataset/__init__.py +++ b/dataset/__init__.py @@ -7,9 +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 +from sqlalchemy import Integer, Text -__all__ = ['Database', 'Table', 'freeze', 'connect', 'Integer', 'String'] +__all__ = ['Database', 'Table', 'freeze', 'connect'] def connect(url=None, schema=None, reflectMetadata=True): diff --git a/dataset/persistence/database.py b/dataset/persistence/database.py index 1ae27b2..bf6d09a 100644 --- a/dataset/persistence/database.py +++ b/dataset/persistence/database.py @@ -7,7 +7,7 @@ from migrate.versioning.util import construct_engine from sqlalchemy.pool import NullPool from sqlalchemy.schema import MetaData, Column, Index from sqlalchemy.schema import Table as SQLATable -from sqlalchemy import Integer, String +from sqlalchemy import Integer, Text from dataset.persistence.table import Table from dataset.persistence.util import ResultIter @@ -97,13 +97,13 @@ class Database(object): return list(set(self.metadata.tables.keys() + self._tables.keys())) - def create_table(self, table_name, primary_id='id', primary_type=Integer): + 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)) or an Integer as specified with primary_type flag. + Text 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,22 +112,25 @@ 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) + + # custom id and type + table2 = db.create_table('population2', 'age') + table3 = db.create_table('population3', primary_id='race', primary_type='Text') """ self._acquire() try: log.debug("Creating table: %s on %r" % (table_name, self.engine)) table = SQLATable(table_name, self.metadata) - if primary_type is Integer: + 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) + elif primary_type is 'Text': + col = Column(primary_id, Text, primary_key=True) else: - raise DatasetException("The primary_type has to be either int or str.") + raise DatasetException( + "The primary_type has to be either 'Integer' or 'Text'.") table.append_column(col) table.create(self.engine) @@ -157,10 +160,12 @@ class Database(object): finally: self._release() - def get_table(self, table_name): + def get_table(self, table_name, primary_id='id', primary_type='Integer'): """ Smart wrapper around *load_table* and *create_table*. Either loads a table or creates it if it doesn't exist yet. + For short-hand to create a table with custom id and type using [], where + table_name, primary_id, and primary_type are specified as a tuple Returns a :py:class:`Table ` instance. :: @@ -168,6 +173,10 @@ class Database(object): table = db.get_table('population') # you can also use the short-hand syntax: table = db['population'] + + # custom id and type + table2 = db['population2', 'age'] # default type is 'Integer' + table3 = db['population3', 'race', 'Text'] """ if table_name in self._tables: return Table(self, self._tables[table_name]) @@ -176,12 +185,15 @@ class Database(object): if self.engine.has_table(table_name, schema=self.schema): return self.load_table(table_name) else: - return self.create_table(table_name) + return self.create_table(table_name, primary_id, primary_type) finally: self._release() def __getitem__(self, table_name): - return self.get_table(table_name) + if type(table_name) is tuple: + return self.get_table(*table_name[:3]) + else: + return self.get_table(table_name) def query(self, query, **kw): """ diff --git a/test/test_persistence.py b/test/test_persistence.py index 063032c..209dbde 100644 --- a/test/test_persistence.py +++ b/test/test_persistence.py @@ -2,7 +2,7 @@ import os import unittest from datetime import datetime -from dataset import connect, Integer, String +from dataset import connect 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, primary_type=String) + table = self.db.create_table("foo2", pid, 'Text') 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, primary_type=Integer) + table = self.db.create_table("foo3", primary_id = pid) assert table.table.exists() assert len(table.table.columns) == 1, table.table.columns assert pid in table.table.c, table.table.c @@ -58,6 +58,31 @@ class DatabaseTestCase(unittest.TestCase): with self.assertRaises(IntegrityError): table.insert({'int_id': 123}) + def test_create_table_shorthand1(self): + pid = "int_id" + table = self.db['foo4', pid] + assert table.table.exists + assert len(table.table.columns) == 1, table.table.columns + assert pid in table.table.c, table.table.c + + table.insert({'int_id': 123}) + table.insert({'int_id': 124}) + assert table.find_one(int_id = 123)[0] == 123 + assert table.find_one(int_id = 124)[0] == 124 + with self.assertRaises(IntegrityError): + table.insert({'int_id': 123}) + + def test_create_table_shorthand2(self): + pid = "string_id" + table = self.db['foo5', pid, 'Text'] + assert table.table.exists + assert len(table.table.columns) == 1, table.table.columns + assert pid in table.table.c, table.table.c + + table.insert({ + 'string_id': 'foobar'}) + assert table.find_one(string_id = 'foobar')[0] == 'foobar' + def test_load_table(self): tbl = self.db.load_table('weather') assert tbl.table == self.tbl.table