diff --git a/pgdata/__init__.py b/pgdata/__init__.py index 7527d69..7c9bdaf 100644 --- a/pgdata/__init__.py +++ b/pgdata/__init__.py @@ -53,6 +53,8 @@ def drop_db(url): FROM pg_database WHERE datname = '{db_name}'""".format(db_name=db_name) if db.query(q).fetchone(): + # DROP DATABASE must be run outside of a transaction conn = db.engine.connect() + conn.execute("commit") conn.execute("DROP DATABASE "+db_name) conn.close() diff --git a/tests/test_pg2ogr.py b/tests/test_pg2ogr.py index dfe7f17..fa36a61 100644 --- a/tests/test_pg2ogr.py +++ b/tests/test_pg2ogr.py @@ -5,13 +5,13 @@ import fiona -from pgdata import connect +import pgdata URL = 'postgresql://postgres:postgres@localhost:5432/pgdata' -DB = connect(URL) +DB = pgdata.connect(URL) DB.execute('CREATE SCHEMA IF NOT EXISTS pgdata') -DB1 = connect(URL) + DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data') AIRPORTS = os.path.join(DATA, 'bc_airports.json') @@ -21,11 +21,8 @@ class ogrpg(unittest.TestCase): def setUp(self): self.tempdir = tempfile.mkdtemp() - def tearDown(self): - shutil.rmtree(self.tempdir) - def test_ogr2pg(self): - db = DB1 + db = DB db.ogr2pg(AIRPORTS, in_layer='bc_airports', out_layer='bc_airports', schema='pgdata') airports = db['pgdata.bc_airports'] @@ -33,14 +30,14 @@ def test_ogr2pg(self): assert sum(1 for _ in airports.all()) == 425 def test_pg2geojson(self): - db = DB1 + db = DB db.pg2ogr(sql='SELECT * FROM pgdata.bc_airports LIMIT 10', driver='GeoJSON', outfile=os.path.join(self.tempdir, 'test_dump.json')) c = fiona.open(os.path.join(self.tempdir, 'test_dump.json'), 'r') assert len(c) == 10 def test_pg2gpkg(self): - db = DB1 + db = DB db.pg2ogr(sql='SELECT * FROM pgdata.bc_airports LIMIT 10', driver='GPKG', outfile=os.path.join(self.tempdir, 'test_dump.gpkg'), outlayer='bc_airports') @@ -48,7 +45,7 @@ def test_pg2gpkg(self): assert len(c) == 10 def test_pg2gpkg_update(self): - db = DB1 + db = DB db.pg2ogr(sql='SELECT * FROM pgdata.bc_airports LIMIT 10', driver='GPKG', outfile=os.path.join(self.tempdir, 'test_dump.gpkg'), outlayer='bc_airports') @@ -59,7 +56,7 @@ def test_pg2gpkg_update(self): assert len(layers) == 2 def test_pg2ogr_append(self): - db = DB1 + db = DB db.pg2ogr(sql='SELECT * FROM pgdata.bc_airports LIMIT 10', driver='GPKG', outfile=os.path.join(self.tempdir, 'test_dump.gpkg'), outlayer='bc_airports') @@ -68,3 +65,10 @@ def test_pg2ogr_append(self): outlayer='bc_airports', append=True) c = fiona.open(os.path.join(self.tempdir, 'test_dump.gpkg'), 'r') assert len(c) == 20 + + def tearDown(self): + shutil.rmtree(self.tempdir) + + +def test_tearDown(): + DB.drop_schema('pgdata', cascade=True)