From fdc2deb2fd38366758b05b9aee30fd8db525fc0d Mon Sep 17 00:00:00 2001 From: Ryan Leckey Date: Mon, 4 Nov 2013 13:35:47 -0800 Subject: [PATCH 1/5] Add database functions; create and drop DB. --- setup.py | 1 + sqlalchemy_utils/__init__.py | 8 ++- sqlalchemy_utils/functions/__init__.py | 6 +- sqlalchemy_utils/functions/database.py | 86 ++++++++++++++++++++++++++ tests/functions/__init__.py | 7 +++ tests/functions/test_database.py | 38 ++++++++++++ 6 files changed, 144 insertions(+), 2 deletions(-) create mode 100644 sqlalchemy_utils/functions/database.py create mode 100644 tests/functions/test_database.py diff --git a/setup.py b/setup.py index afa2336..9cb382e 100644 --- a/setup.py +++ b/setup.py @@ -34,6 +34,7 @@ extras_require = { 'docutils>=0.10', 'flexmock>=0.9.7', 'psycopg2>=2.4.6', + 'pymysql' ], 'anyjson': ['anyjson>=0.3.3'], 'babel': ['Babel>=1.3'], diff --git a/sqlalchemy_utils/__init__.py b/sqlalchemy_utils/__init__.py index dcc636d..7911b77 100644 --- a/sqlalchemy_utils/__init__.py +++ b/sqlalchemy_utils/__init__.py @@ -13,7 +13,10 @@ from .functions import ( mock_engine, sort_query, table_name, - with_backrefs + with_backrefs, + database_exists, + create_database, + drop_database ) from .listeners import coercion_listener from .merge import merge, Merger @@ -96,4 +99,7 @@ __all__ = ( TimezoneType, TSVectorType, UUIDType, + database_exists, + create_database, + drop_database ) diff --git a/sqlalchemy_utils/functions/__init__.py b/sqlalchemy_utils/functions/__init__.py index 0ac53be..9cce2a0 100644 --- a/sqlalchemy_utils/functions/__init__.py +++ b/sqlalchemy_utils/functions/__init__.py @@ -6,6 +6,7 @@ from .defer_except import defer_except from .mock import create_mock_engine, mock_engine from .render import render_expression, render_statement from .sort_query import sort_query, QuerySorterException +from .database import database_exists, create_database, drop_database __all__ = ( @@ -18,7 +19,10 @@ __all__ = ( render_statement, with_backrefs, CompositePath, - QuerySorterException + QuerySorterException, + database_exists, + create_database, + drop_database ) diff --git a/sqlalchemy_utils/functions/database.py b/sqlalchemy_utils/functions/database.py new file mode 100644 index 0000000..be2ebe6 --- /dev/null +++ b/sqlalchemy_utils/functions/database.py @@ -0,0 +1,86 @@ +from sqlalchemy.engine.url import make_url +import sqlalchemy as sa +from sqlalchemy.exc import ProgrammingError +import os + + +def database_exists(url): + """Check if a database exists. + """ + + url = make_url(url) + database = url.database + url.database = None + + engine = sa.create_engine(url) + + if engine.dialect.name == 'postgres': + text = "SELECT 1 FROM pg_database WHERE datname='%s'" % database + return bool(engine.execute(text).scalar()) + + elif engine.dialect.name == 'mysql': + text = ("SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA " + "WHERE SCHEMA_NAME = '%s'" % database) + return bool(engine.execute(text).scalar()) + + elif engine.dialect.name == 'sqlite': + return database == ':memory:' or os.path.exists(database) + + else: + text = 'SELECT 1' + try: + url.database = database + engine = sa.create_engine(url) + engine.execute(text) + return True + + except ProgrammingError: + return False + + +def create_database(url, encoding='utf8'): + """Issue the appropriate CREATE DATABASE statement. + """ + + url = make_url(url) + + database = url.database + if not url.drivername.startswith('sqlite'): + url.database = None + + engine = sa.create_engine(url) + + if engine.dialect.name == 'postgres': + text = "CREATE DATABASE %s ENCODING = '%s'" % (database, encoding) + engine.execute(text) + + elif engine.dialect.name == 'mysql': + text = "CREATE DATABASE %s CHARACTER SET = '%s'" % (database, encoding) + engine.execute(text) + + elif engine.dialect.name == 'sqlite' and database != ':memory:': + open(database, 'w').close() + + else: + text = "CREATE DATABASE %s" % database + engine.execute(text) + + +def drop_database(url): + """Issue the appropriate DROP DATABASE statement. + """ + + url = make_url(url) + + database = url.database + if not url.drivername.startswith('sqlite'): + url.database = None + + engine = sa.create_engine(url) + + if engine.dialect.name == 'sqlite' and url.database != ':memory:': + os.remove(url.database) + + else: + text = "DROP DATABASE %s" % database + engine.execute(text) diff --git a/tests/functions/__init__.py b/tests/functions/__init__.py index e69de29..ad6e251 100644 --- a/tests/functions/__init__.py +++ b/tests/functions/__init__.py @@ -0,0 +1,7 @@ +import sqlalchemy as sa +from tests import TestCase +from sqlalchemy_utils.functions import ( + render_statement, + render_expression, + mock_engine +) diff --git a/tests/functions/test_database.py b/tests/functions/test_database.py new file mode 100644 index 0000000..34559b1 --- /dev/null +++ b/tests/functions/test_database.py @@ -0,0 +1,38 @@ +import sqlalchemy as sa +import os +from tests import TestCase +from sqlalchemy_utils import ( + create_database, + drop_database, + database_exists, +) + + +class DatabaseTest(TestCase): + + def test_create_and_drop(self): + assert not database_exists(self.url) + create_database(self.url) + assert database_exists(self.url) + drop_database(self.url) + assert not database_exists(self.url) + + +class TestDatabaseSQLite(DatabaseTest): + + url = 'sqlite:///sqlalchemy_utils.db' + + def setup(self): + if os.path.exists('sqlalchemy_utils.db'): + os.remove('sqlalchemy_utils.db') + + def test_exists_memory(self): + assert database_exists('sqlite:///:memory:') + + +class TestDatabaseMySQL(DatabaseTest): + url = 'mysql+pymysql://travis@localhost/db_test_sqlalchemy_util' + + +class TestDatabasePostgres(DatabaseTest): + url = 'postgres://postgres@localhost/db_test_sqlalchemy_util' From 178993d9f1da3b1bfade0b3fca076bd069936115 Mon Sep 17 00:00:00 2001 From: Ryan Leckey Date: Mon, 4 Nov 2013 13:54:48 -0800 Subject: [PATCH 2/5] Fix name of postgresql dialect. --- sqlalchemy_utils/functions/database.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sqlalchemy_utils/functions/database.py b/sqlalchemy_utils/functions/database.py index be2ebe6..7871e4b 100644 --- a/sqlalchemy_utils/functions/database.py +++ b/sqlalchemy_utils/functions/database.py @@ -1,6 +1,6 @@ from sqlalchemy.engine.url import make_url import sqlalchemy as sa -from sqlalchemy.exc import ProgrammingError +from sqlalchemy.exc import ProgrammingError, OperationalError import os @@ -14,7 +14,7 @@ def database_exists(url): engine = sa.create_engine(url) - if engine.dialect.name == 'postgres': + if engine.dialect.name == 'postgresql': text = "SELECT 1 FROM pg_database WHERE datname='%s'" % database return bool(engine.execute(text).scalar()) @@ -34,7 +34,7 @@ def database_exists(url): engine.execute(text) return True - except ProgrammingError: + except (ProgrammingError, OperationalError): return False @@ -50,7 +50,7 @@ def create_database(url, encoding='utf8'): engine = sa.create_engine(url) - if engine.dialect.name == 'postgres': + if engine.dialect.name == 'postgresql': text = "CREATE DATABASE %s ENCODING = '%s'" % (database, encoding) engine.execute(text) From 7e14156379753118f94b220b4c02cf2e9d5479f2 Mon Sep 17 00:00:00 2001 From: Ryan Leckey Date: Mon, 4 Nov 2013 14:08:57 -0800 Subject: [PATCH 3/5] Ensure url is unchanged. --- sqlalchemy_utils/functions/database.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/sqlalchemy_utils/functions/database.py b/sqlalchemy_utils/functions/database.py index 7871e4b..890eace 100644 --- a/sqlalchemy_utils/functions/database.py +++ b/sqlalchemy_utils/functions/database.py @@ -2,13 +2,14 @@ from sqlalchemy.engine.url import make_url import sqlalchemy as sa from sqlalchemy.exc import ProgrammingError, OperationalError import os +from copy import copy def database_exists(url): """Check if a database exists. """ - url = make_url(url) + url = copy(url) database = url.database url.database = None @@ -42,7 +43,7 @@ def create_database(url, encoding='utf8'): """Issue the appropriate CREATE DATABASE statement. """ - url = make_url(url) + url = copy(url) database = url.database if not url.drivername.startswith('sqlite'): @@ -70,7 +71,7 @@ def drop_database(url): """Issue the appropriate DROP DATABASE statement. """ - url = make_url(url) + url = copy(url) database = url.database if not url.drivername.startswith('sqlite'): From eb53f7b31ea994cd17665ff31aef38f96d12cbde Mon Sep 17 00:00:00 2001 From: Ryan Leckey Date: Mon, 4 Nov 2013 14:26:46 -0800 Subject: [PATCH 4/5] Ensure isolation level is autocommit for postgres. --- sqlalchemy_utils/functions/database.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/sqlalchemy_utils/functions/database.py b/sqlalchemy_utils/functions/database.py index 890eace..cf94740 100644 --- a/sqlalchemy_utils/functions/database.py +++ b/sqlalchemy_utils/functions/database.py @@ -52,6 +52,11 @@ def create_database(url, encoding='utf8'): engine = sa.create_engine(url) if engine.dialect.name == 'postgresql': + if engine.driver == 'psycopg2': + from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT + engine.raw_connection().set_isolation_level( + ISOLATION_LEVEL_AUTOCOMMIT) + text = "CREATE DATABASE %s ENCODING = '%s'" % (database, encoding) engine.execute(text) @@ -82,6 +87,12 @@ def drop_database(url): if engine.dialect.name == 'sqlite' and url.database != ':memory:': os.remove(url.database) + elif engine.dialect.name == 'postgresql' and engine.driver == 'psycopg2': + from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT + engine.raw_connection().set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT) + text = "DROP DATABASE %s" % database + engine.execute(text) + else: text = "DROP DATABASE %s" % database engine.execute(text) From 5f3266541e00ccb21cd020a698f988b459578e56 Mon Sep 17 00:00:00 2001 From: Ryan Leckey Date: Mon, 4 Nov 2013 14:48:46 -0800 Subject: [PATCH 5/5] Ensure we get URL objects. --- sqlalchemy_utils/functions/database.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sqlalchemy_utils/functions/database.py b/sqlalchemy_utils/functions/database.py index cf94740..8c16e71 100644 --- a/sqlalchemy_utils/functions/database.py +++ b/sqlalchemy_utils/functions/database.py @@ -9,7 +9,7 @@ def database_exists(url): """Check if a database exists. """ - url = copy(url) + url = copy(make_url(url)) database = url.database url.database = None @@ -43,7 +43,7 @@ def create_database(url, encoding='utf8'): """Issue the appropriate CREATE DATABASE statement. """ - url = copy(url) + url = copy(make_url(url)) database = url.database if not url.drivername.startswith('sqlite'): @@ -76,7 +76,7 @@ def drop_database(url): """Issue the appropriate DROP DATABASE statement. """ - url = copy(url) + url = copy(make_url(url)) database = url.database if not url.drivername.startswith('sqlite'):