connection.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209
  1. # Licensed to the Apache Software Foundation (ASF) under one or more
  2. # contributor license agreements. See the NOTICE file distributed with
  3. # this work for additional information regarding copyright ownership.
  4. # The ASF licenses this file to You under the Apache License, Version 2.0
  5. # (the "License"); you may not use this file except in compliance with
  6. # the License. You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. import logging
  16. import uuid
  17. import weakref
  18. from phoenixdb import errors
  19. from phoenixdb.cursor import Cursor
  20. from phoenixdb.errors import ProgrammingError
  21. from phoenixdb.meta import Meta
  22. __all__ = ['Connection']
  23. logger = logging.getLogger(__name__)
  24. AVATICA_PROPERTIES = ('autoCommit', 'autocommit', 'readOnly', 'readonly', 'transactionIsolation',
  25. 'catalog', 'schema')
  26. class Connection(object):
  27. """Database connection.
  28. You should not construct this object manually, use :func:`~phoenixdb.connect` instead.
  29. """
  30. cursor_factory = None
  31. """
  32. The default cursor factory used by :meth:`cursor` if the parameter is not specified.
  33. """
  34. def __init__(self, client, cursor_factory=None, **kwargs):
  35. self._client = client
  36. self._closed = False
  37. if cursor_factory is not None:
  38. self.cursor_factory = cursor_factory
  39. else:
  40. self.cursor_factory = Cursor
  41. self._cursors = []
  42. self._phoenix_props, avatica_props_init = Connection._map_conn_props(kwargs)
  43. self.open()
  44. # TODO we could probably optimize it away if the defaults are not changed
  45. self.set_session(**avatica_props_init)
  46. def __del__(self):
  47. if not self._closed:
  48. self.close()
  49. def __enter__(self):
  50. return self
  51. def __exit__(self, exc_type, exc_value, traceback):
  52. if not self._closed:
  53. self.close()
  54. @staticmethod
  55. def _default_avatica_props():
  56. return {'autoCommit': False,
  57. 'readOnly': False,
  58. 'transactionIsolation': 0,
  59. 'catalog': '',
  60. 'schema': ''}
  61. @staticmethod
  62. def _map_conn_props(conn_props):
  63. """Sorts and prepocesses args that should be passed to Phoenix and Avatica"""
  64. avatica_props = dict([(k, conn_props[k]) for k in conn_props.keys() if k in AVATICA_PROPERTIES])
  65. phoenix_props = dict([(k, conn_props[k]) for k in conn_props.keys() if k not in AVATICA_PROPERTIES])
  66. avatica_props = Connection._map_legacy_avatica_props(avatica_props)
  67. return (phoenix_props, avatica_props)
  68. @staticmethod
  69. def _map_legacy_avatica_props(props):
  70. if 'autocommit' in props:
  71. props['autoCommit'] = bool(props.pop('autocommit'))
  72. if 'readonly' in props:
  73. props['readOnly'] = bool(props.pop('readonly'))
  74. return props
  75. def open(self):
  76. """Opens the connection."""
  77. self._id = str(uuid.uuid4())
  78. self._client.open_connection(self._id, info=self._phoenix_props)
  79. def close(self):
  80. """Closes the connection.
  81. No further operations are allowed, either on the connection or any
  82. of its cursors, once the connection is closed.
  83. If the connection is used in a ``with`` statement, this method will
  84. be automatically called at the end of the ``with`` block.
  85. """
  86. if self._closed:
  87. raise ProgrammingError('The connection is already closed.')
  88. for cursor_ref in self._cursors:
  89. cursor = cursor_ref()
  90. if cursor is not None and not cursor._closed:
  91. cursor.close()
  92. self._client.close_connection(self._id)
  93. self._client.close()
  94. self._closed = True
  95. @property
  96. def closed(self):
  97. """Read-only attribute specifying if the connection is closed or not."""
  98. return self._closed
  99. def commit(self):
  100. if self._closed:
  101. raise ProgrammingError('The connection is already closed.')
  102. self._client.commit(self._id)
  103. def rollback(self):
  104. if self._closed:
  105. raise ProgrammingError('The connection is already closed.')
  106. self._client.rollback(self._id)
  107. def cursor(self, cursor_factory=None):
  108. """Creates a new cursor.
  109. :param cursor_factory:
  110. This argument can be used to create non-standard cursors.
  111. The class returned must be a subclass of
  112. :class:`~phoenixdb.cursor.Cursor` (for example :class:`~phoenixdb.cursor.DictCursor`).
  113. A default factory for the connection can also be specified using the
  114. :attr:`cursor_factory` attribute.
  115. :returns:
  116. A :class:`~phoenixdb.cursor.Cursor` object.
  117. """
  118. if self._closed:
  119. raise ProgrammingError('The connection is already closed.')
  120. cursor = (cursor_factory or self.cursor_factory)(self)
  121. self._cursors.append(weakref.ref(cursor, self._cursors.remove))
  122. return cursor
  123. def set_session(self, **props):
  124. """Sets one or more parameters in the current connection.
  125. :param autocommit:
  126. Switch the connection to autocommit mode.
  127. :param readonly:
  128. Switch the connection to read-only mode.
  129. """
  130. props = Connection._map_legacy_avatica_props(props)
  131. self._avatica_props = self._client.connection_sync_dict(self._id, props)
  132. @property
  133. def autocommit(self):
  134. """Read/write attribute for switching the connection's autocommit mode."""
  135. return self._avatica_props['autoCommit']
  136. @autocommit.setter
  137. def autocommit(self, value):
  138. if self._closed:
  139. raise ProgrammingError('The connection is already closed.')
  140. self._avatica_props = self._client.connection_sync_dict(self._id, {'autoCommit': bool(value)})
  141. @property
  142. def readonly(self):
  143. """Read/write attribute for switching the connection's readonly mode."""
  144. return self._avatica_props['readOnly']
  145. @readonly.setter
  146. def readonly(self, value):
  147. if self._closed:
  148. raise ProgrammingError('The connection is already closed.')
  149. self._avatica_props = self._client.connection_sync_dict(self._id, {'readOnly': bool(value)})
  150. @property
  151. def transactionisolation(self):
  152. return self._avatica_props['_transactionIsolation']
  153. @transactionisolation.setter
  154. def transactionisolation(self, value):
  155. if self._closed:
  156. raise ProgrammingError('The connection is already closed.')
  157. self._avatica_props = self._client.connection_sync_dict(self._id, {'transactionIsolation': bool(value)})
  158. def meta(self):
  159. """Creates a new meta.
  160. :returns:
  161. A :class:`~phoenixdb.meta` object.
  162. """
  163. if self._closed:
  164. raise ProgrammingError('The connection is already closed.')
  165. meta = Meta(self)
  166. return meta
  167. for name in errors.__all__:
  168. setattr(Connection, name, getattr(errors, name))