| 1 | import sqlalchemy |
| 2 | |
| 3 | class Connection(object): |
| 4 | current = None |
| 5 | connections = {} |
| 6 | @classmethod |
| 7 | def tell_format(cls): |
| 8 | return "Format: (postgresql|mysql)://username:password@hostname/dbname, or one of %s" \ |
| 9 | % str(cls.connections.keys()) |
| 10 | def __init__(self, connect_str=None): |
| 11 | try: |
| 12 | engine = sqlalchemy.create_engine(connect_str) |
| 13 | except: # TODO: bare except; but what's an ArgumentError? |
| 14 | print(self.tell_format()) |
| 15 | raise |
| 16 | self.dialect = engine.url.get_dialect() |
| 17 | self.metadata = sqlalchemy.MetaData(bind=engine) |
| 18 | self.name = self.assign_name(engine) |
| 19 | self.session = engine.connect() |
| 20 | self.connections[self.name] = self |
| 21 | self.connections[str(self.metadata.bind.url)] = self |
| 22 | Connection.current = self |
| 23 | @classmethod |
| 24 | def get(cls, descriptor): |
| 25 | if isinstance(descriptor, Connection): |
| 26 | cls.current = descriptor |
| 27 | elif descriptor: |
| 28 | conn = cls.connections.get(descriptor) or \ |
| 29 | cls.connections.get(descriptor.lower()) |
| 30 | if conn: |
| 31 | cls.current = conn |
| 32 | else: |
| 33 | cls.current = Connection(descriptor) |
| 34 | if cls.current: |
| 35 | return cls.current |
| 36 | else: |
| 37 | raise Exception(cls.tell_format()) |
| 38 | @classmethod |
| 39 | def assign_name(cls, engine): |
| 40 | core_name = '%s@%s' % (engine.url.username, engine.url.database) |
| 41 | incrementer = 1 |
| 42 | name = core_name |
| 43 | while name in cls.connections: |
| 44 | name = '%s_%d' % (core_name, incrementer) |
| 45 | incrementer += 1 |
| 46 | return name |