Selection of existing connection works with concealed password

pull/116/merge v0.3.9
Catherine Devlin 2018-04-05 22:14:12 -04:00
parent 59a9719c26
commit 1e8161548f
3 changed files with 21 additions and 11 deletions

View File

@ -149,3 +149,4 @@ Deleted Plugin import left behind in 0.2.2
* added README example (thanks tanhuil)
* bugfix in executing column_local_vars (thanks tebeka)
* pgspecial installation optional (thanks jstoebel and arjoe)
* conceal passwords in connection strings (thanks jstoebel)

View File

@ -1,3 +1,3 @@
#!/bin/bash
python -c "import pytest; pytest.main(['.', '-x', '--pdb'])"
ipython -c "import pytest; pytest.main(['.', '-x', '--pdb'])"
# Insert breakpoints with `import pytest; pytest.set_trace()`

View File

@ -1,10 +1,26 @@
import sqlalchemy
import os
import re
class ConnectionError(Exception):
pass
def rough_dict_get(dct, sought, default=None):
'''
Like dct.get(sought), but any key containing sought will do.
If there is a `@` in sought, seek each piece separately.
This lets `me@server` match `me:***@myserver/db`
'''
sought = sought.split('@')
for (key, val) in dct.items():
if not any(s.lower() not in key.lower() for s in sought):
return val
return default
class Connection(object):
current = None
connections = {}
@ -25,8 +41,7 @@ class Connection(object):
self.metadata = sqlalchemy.MetaData(bind=engine)
self.name = self.assign_name(engine)
self.session = engine.connect()
self.connections[self.name] = self
self.connections[str(self.metadata.bind.url)] = self
self.connections[repr(self.metadata.bind.url)] = self
Connection.current = self
@classmethod
@ -37,8 +52,7 @@ class Connection(object):
if isinstance(descriptor, Connection):
cls.current = descriptor
else:
existing = cls.connections.get(descriptor) or \
cls.connections.get(descriptor.lower())
existing = rough_dict_get(cls.connections, descriptor)
cls.current = existing or Connection(descriptor)
else:
if cls.connections:
@ -52,12 +66,7 @@ class Connection(object):
@classmethod
def assign_name(cls, engine):
core_name = '%s@%s' % (engine.url.username or '', engine.url.database)
incrementer = 1
name = core_name
while name in cls.connections:
name = '%s_%d' % (core_name, incrementer)
incrementer += 1
name = '%s@%s' % (engine.url.username or '', engine.url.database)
return name
@classmethod