Skip to content

Commit

Permalink
Fixed a bug when querying with schema path to Redshift/Postgresql
Browse files Browse the repository at this point in the history
  • Loading branch information
sungjuly committed Feb 26, 2017
1 parent 10ea635 commit bca6330
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 7 deletions.
24 changes: 17 additions & 7 deletions superset/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -776,15 +776,25 @@ def get_sqla_engine(self, schema=None):
extra = self.get_extra()
url = make_url(self.sqlalchemy_uri_decrypted)
params = extra.get('engine_params', {})
if self.backend == 'presto' and schema:
if '/' in url.database:
url.database = url.database.split('/')[0] + '/' + schema
else:
url.database += '/' + schema
elif schema:
url.database = schema
url.database = self.get_database_for_various_backend(url, schema)
return create_engine(url, **params)

def get_database_for_various_backend(self, uri, default_database=None):
database = uri.database
if self.backend == 'presto' and default_database:
if '/' in database:
database = database.split('/')[0] + '/' + default_database
else:
database += '/' + default_database
# Postgres and Redshift use the concept of schema as a logical entity
# on top of the database, so the database should not be changed
# even if passed default_database
elif self.backend == 'redshift' or self.backend == 'postgresql':
pass
elif default_database:
database = default_database
return database

def get_reserved_words(self):
return self.get_sqla_engine().dialect.preparer.reserved_words

Expand Down
48 changes: 48 additions & 0 deletions tests/model_tests.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
import unittest

from sqlalchemy.engine.url import make_url

from superset.models import Database


class DatabaseModelTestCase(unittest.TestCase):
def test_database_for_various_backend(self):
sqlalchemy_uri = 'presto://presto.airbnb.io:8080/hive/default'
model = Database(sqlalchemy_uri=sqlalchemy_uri)
url = make_url(model.sqlalchemy_uri)
db = model.get_database_for_various_backend(url, None)
assert db == 'hive/default'
db = model.get_database_for_various_backend(url, 'raw_data')
assert db == 'hive/raw_data'

sqlalchemy_uri = 'redshift+psycopg2://superset:[email protected]:5439/prod'
model = Database(sqlalchemy_uri=sqlalchemy_uri)
url = make_url(model.sqlalchemy_uri)
db = model.get_database_for_various_backend(url, None)
assert db == 'prod'
db = model.get_database_for_various_backend(url, 'test')
assert db == 'prod'

sqlalchemy_uri = 'postgresql+psycopg2://superset:[email protected]:5439/prod'
model = Database(sqlalchemy_uri=sqlalchemy_uri)
url = make_url(model.sqlalchemy_uri)
db = model.get_database_for_various_backend(url, None)
assert db == 'prod'
db = model.get_database_for_various_backend(url, 'adhoc')
assert db == 'prod'

sqlalchemy_uri = 'hive://[email protected]:10000/raw_data'
model = Database(sqlalchemy_uri=sqlalchemy_uri)
url = make_url(model.sqlalchemy_uri)
db = model.get_database_for_various_backend(url, None)
assert db == 'raw_data'
db = model.get_database_for_various_backend(url, 'adhoc')
assert db == 'adhoc'

sqlalchemy_uri = 'mysql://superset:[email protected]/superset'
model = Database(sqlalchemy_uri=sqlalchemy_uri)
url = make_url(model.sqlalchemy_uri)
db = model.get_database_for_various_backend(url, None)
assert db == 'superset'
db = model.get_database_for_various_backend(url, 'adhoc')
assert db == 'adhoc'

0 comments on commit bca6330

Please sign in to comment.