forked from catherinedevlin/ipython-sql
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathconnection.py
More file actions
122 lines (104 loc) · 3.92 KB
/
Copy pathconnection.py
File metadata and controls
122 lines (104 loc) · 3.92 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import os
import re
import sqlalchemy
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 = {}
@classmethod
def tell_format(cls):
return """Connection info needed in SQLAlchemy format, example:
postgresql://username:password@hostname/dbname
or an existing connection: %s""" % str(
cls.connections.keys()
)
def __init__(self, connect_str=None, connect_args={}, creator=None):
try:
if creator:
engine = sqlalchemy.create_engine(
connect_str, connect_args=connect_args, creator=creator
)
else:
engine = sqlalchemy.create_engine(
connect_str, connect_args=connect_args
)
except: # TODO: bare except; but what's an ArgumentError?
print(self.tell_format())
raise
self.dialect = engine.url.get_dialect()
self.metadata = sqlalchemy.MetaData(bind=engine)
self.name = self.assign_name(engine)
self.session = engine.connect()
self.connections[repr(self.metadata.bind.url)] = self
self.connect_args = connect_args
Connection.current = self
@classmethod
def set(cls, descriptor, displaycon, connect_args={}, creator=None):
"Sets the current database connection"
if descriptor:
if isinstance(descriptor, Connection):
cls.current = descriptor
else:
existing = rough_dict_get(cls.connections, descriptor)
# http://docs.sqlalchemy.org/en/rel_0_9/core/engines.html#custom-dbapi-connect-arguments
cls.current = existing or Connection(descriptor, connect_args, creator)
else:
if cls.connections:
if displaycon:
print(cls.connection_list())
else:
if os.getenv("DATABASE_URL"):
cls.current = Connection(
os.getenv("DATABASE_URL"), connect_args, creator
)
else:
raise ConnectionError(
"Environment variable $DATABASE_URL not set, and no connect string given."
)
return cls.current
@classmethod
def assign_name(cls, engine):
name = "%s@%s" % (engine.url.username or "", engine.url.database)
return name
@classmethod
def connection_list(cls):
result = []
for key in sorted(cls.connections):
engine_url = cls.connections[
key
].metadata.bind.url # type: sqlalchemy.engine.url.URL
if cls.connections[key] == cls.current:
template = " * {}"
else:
template = " {}"
result.append(template.format(engine_url.__repr__()))
return "\n".join(result)
@classmethod
def _close(cls, descriptor):
if isinstance(descriptor, Connection):
conn = descriptor
else:
conn = cls.connections.get(descriptor) or cls.connections.get(
descriptor.lower()
)
if not conn:
raise Exception(
"Could not close connection because it was not found amongst these: %s"
% str(cls.connections.keys())
)
cls.connections.pop(str(conn.metadata.bind.url))
conn.session.close()
def close(self):
self.__class__._close(self)