Fix: DynamoDB test connection was broken

This commit is contained in:
Arik Fraimovich 2016-12-27 14:47:12 +02:00
parent f5e8fc816c
commit 1388e76b48

View File

@ -34,8 +34,6 @@ types_map = {
class DynamoDBSQL(BaseSQLQueryRunner):
noop_query = "SELECT 1"
@classmethod
def configuration_schema(cls):
return {
@ -45,29 +43,21 @@ class DynamoDBSQL(BaseSQLQueryRunner):
"type": "string",
"default": "us-east-1"
},
"host": {
"type": "string",
"default": "Use for non standard endpoints."
},
"port": {
"type": "number",
"default": 80
},
"access_key": {
"type": "string",
},
"secret_key": {
"type": "string",
},
"is_secure": {
"type": "boolean",
"default": False,
}
},
"required": ["access_key", "secret_key"],
"secret": ["secret_key"]
}
def test_connection(self):
engine = self._connect()
list(engine.connection.list_tables())
@classmethod
def annotate_query(cls):
return False
@ -93,24 +83,20 @@ class DynamoDBSQL(BaseSQLQueryRunner):
if config.get('host') == '':
config['host'] = None
return engine, engine.connect(**config)
engine.connect(**config)
return engine
def _get_tables(self, schema):
engine = self._connect()
try:
engine, _ = self._connect()
for table in engine.describe_all():
schema[table.name] = {'name': table.name, 'columns': table.attrs.keys()}
except Exception as e:
logging.exception(e)
raise sys.exc_info()[1], None, sys.exc_info()[2]
for table in engine.describe_all():
schema[table.name] = {'name': table.name, 'columns': table.attrs.keys()}
def run_query(self, query, user):
connection = None
engine = None
try:
engine, connection = self._connect()
engine = self._connect()
res_dict = engine.execute(query if str(query).endswith(';') else str(query)+';')
@ -137,8 +123,8 @@ class DynamoDBSQL(BaseSQLQueryRunner):
error = e.message
json_data = None
except KeyboardInterrupt:
if connection:
connection.cancel()
if engine and engine.connection:
engine.connection.cancel()
error = "Query cancelled by user."
json_data = None
except Exception as e: