Skip to content

Commit 2e28fc2

Browse files
mambocabmambocab
authored
PYTHON-649 (apache#677)
Adds a test to reproduce PYTHON-649 and fixes it. Also adds docs and tests for some existing connection-management code.
1 parent 3bc5adc commit 2e28fc2

5 files changed

Lines changed: 188 additions & 8 deletions

File tree

‎cassandra/cqlengine/connection.py‎

Lines changed: 69 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,13 @@ def __init__(self, name, hosts, consistency=None,
8181
self.cluster_options = cluster_options if cluster_options else {}
8282
self.lazy_connect_lock = threading.RLock()
8383

84+
@classmethod
85+
def from_session(cls, name, session):
86+
instance = cls(name=name, hosts=session.hosts)
87+
instance.cluster, instance.session = session.cluster, session
88+
instance.setup_session()
89+
return instance
90+
8491
def setup(self):
8592
"""Setup the connection"""
8693
global cluster, session
@@ -132,21 +139,67 @@ def handle_lazy_connect(self):
132139
self.setup()
133140

134141

135-
def register_connection(name, hosts, consistency=None, lazy_connect=False,
136-
retry_connect=False, cluster_options=None, default=False):
142+
def register_connection(name, hosts=None, consistency=None, lazy_connect=False,
143+
retry_connect=False, cluster_options=None, default=False,
144+
session=None):
145+
"""
146+
Add a connection to the connection registry. ``hosts`` and ``session`` are
147+
mutually exclusive, and ``consistency``, ``lazy_connect``,
148+
``retry_connect``, and ``cluster_options`` only work with ``hosts``. Using
149+
``hosts`` will create a new :class:`cassandra.cluster.Cluster` and
150+
:class:`cassandra.cluster.Session`.
151+
152+
:param list hosts: list of hosts, (``contact_points`` for :class:`cassandra.cluster.Cluster`).
153+
:param int consistency: The default :class:`~.ConsistencyLevel` for the
154+
registered connection's new session. Default is the same as
155+
:attr:`.Session.default_consistency_level`. For use with ``hosts`` only;
156+
will fail when used with ``session``.
157+
:param bool lazy_connect: True if should not connect until first use. For
158+
use with ``hosts`` only; will fail when used with ``session``.
159+
:param bool retry_connect: True if we should retry to connect even if there
160+
was a connection failure initially. For use with ``hosts`` only; will
161+
fail when used with ``session``.
162+
:param dict cluster_options: A dict of options to be used as keyword
163+
arguments to :class:`cassandra.cluster.Cluster`. For use with ``hosts``
164+
only; will fail when used with ``session``.
165+
:param bool default: If True, set the new connection as the cqlengine
166+
default
167+
:param Session session: A :class:`cassandra.cluster.Session` to be used in
168+
the created connection.
169+
"""
137170

138171
if name in _connections:
139172
log.warning("Registering connection '{0}' when it already exists.".format(name))
140173

141-
conn = Connection(name, hosts, consistency=consistency,lazy_connect=lazy_connect,
142-
retry_connect=retry_connect, cluster_options=cluster_options)
174+
hosts_xor_session_passed = (hosts is None) ^ (session is None)
175+
if not hosts_xor_session_passed:
176+
raise CQLEngineException(
177+
"Must pass exactly one of 'hosts' or 'session' arguments"
178+
)
179+
elif session is not None:
180+
invalid_config_args = (consistency is not None or
181+
lazy_connect is not False or
182+
retry_connect is not False or
183+
cluster_options is not None)
184+
if invalid_config_args:
185+
raise CQLEngineException(
186+
"Session configuration arguments and 'session' argument are mutually exclusive"
187+
)
188+
conn = Connection.from_session(name, session=session)
189+
conn.setup_session()
190+
elif hosts is not None:
191+
conn = Connection(
192+
name, hosts=hosts,
193+
consistency=consistency, lazy_connect=lazy_connect,
194+
retry_connect=retry_connect, cluster_options=cluster_options
195+
)
196+
conn.setup()
143197

144198
_connections[name] = conn
145199

146200
if default:
147201
set_default_connection(name)
148202

149-
conn.setup()
150203
return conn
151204

152205

@@ -222,7 +275,12 @@ def set_session(s):
222275
This may be relaxed in the future
223276
"""
224277

225-
conn = get_connection()
278+
try:
279+
conn = get_connection()
280+
except CQLEngineException:
281+
# no default connection set; initalize one
282+
register_connection('default', session=s, default=True)
283+
conn = get_connection()
226284

227285
if conn.session:
228286
log.warning("configuring new default connection for cqlengine when one was already set")
@@ -304,7 +362,11 @@ def get_cluster(connection=None):
304362
def register_udt(keyspace, type_name, klass, connection=None):
305363
udt_by_keyspace[keyspace][type_name] = klass
306364

307-
cluster = get_cluster(connection)
365+
try:
366+
cluster = get_cluster(connection)
367+
except CQLEngineException:
368+
cluster = None
369+
308370
if cluster:
309371
try:
310372
cluster.register_user_type(keyspace, type_name, klass)

‎docs/cqlengine/connections.rst‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ Connections are experimental and aimed to ease the use of multiple sessions with
88
Register a new connection
99
=========================
1010

11-
To use cqlengine, you need at least a default connection. This is currently done automatically under the hood with :func:`connection.setup <.connection.setup>`. If you want to use another cluster/session, you need to register a new cqlengine connection. You register a connection with :func:`~.connection.register_connection`
11+
To use cqlengine, you need at least a default connection. If you initialize cqlengine's connections with with :func:`connection.setup <.connection.setup>`, a connection will be created automatically. If you want to use another cluster/session, you need to register a new cqlengine connection. You register a connection with :func:`~.connection.register_connection`:
1212

1313
.. code-block:: python
1414
@@ -17,6 +17,17 @@ To use cqlengine, you need at least a default connection. This is currently done
1717
connection.setup(['127.0.0.1')
1818
connection.register_connection('cluster2', ['127.0.0.2'])
1919
20+
:func:`~.connection.register_connection` can take a list of hosts, as shown above, in which case it will create a connection with a new session. It can also take a `session` argument if you've already created a session:
21+
22+
.. code-block:: python
23+
24+
from cassandra.cqlengine import connection
25+
from cassandra.cluster import Cluster
26+
27+
session = Cluster(['127.0.0.1']).connect()
28+
connection.register_connection('cluster3', session=session)
29+
30+
2031
Change the default connection
2132
=============================
2233

‎tests/integration/cqlengine/test_connections.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
from cassandra import InvalidRequest
16+
from cassandra.cluster import Cluster
1617
from cassandra.cluster import NoHostAvailable
1718
from cassandra.cqlengine import columns, CQLEngineException
1819
from cassandra.cqlengine import connection as conn
@@ -217,6 +218,12 @@ def test_create_drop_table(self):
217218
for ks in self.keyspaces:
218219
drop_keyspace(ks, connections=self.conns)
219220

221+
def test_connection_creation_from_session(self):
222+
session = Cluster(['127.0.0.1']).connect()
223+
connection_name = 'from_session'
224+
conn.register_connection(connection_name, session=session)
225+
self.addCleanup(conn.unregister_connection, connection_name)
226+
220227

221228
class BatchQueryConnectionTests(BaseCassEngTestCase):
222229

‎tests/unit/cqlengine/test_connection.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# Copyright 2013-2016 DataStax, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
try:
16+
import unittest2 as unittest
17+
except ImportError:
18+
import unittest # noqa
19+
20+
from cassandra.cqlengine import connection
21+
from cassandra.query import dict_factory
22+
23+
from mock import Mock
24+
25+
26+
class ConnectionTest(unittest.TestCase):
27+
28+
no_registered_connection_msg = "doesn't exist in the registry"
29+
30+
def setUp(self):
31+
super(ConnectionTest, self).setUp()
32+
self.assertFalse(
33+
connection._connections,
34+
'Test precondition not met: connections are registered: {cs}'.format(cs=connection._connections)
35+
)
36+
37+
def test_set_session_without_existing_connection(self):
38+
"""
39+
Users can set the default session without having a default connection set.
40+
"""
41+
mock_session = Mock(
42+
row_factory=dict_factory,
43+
encoder=Mock(mapping={})
44+
)
45+
connection.set_session(mock_session)
46+
47+
def test_get_session_fails_without_existing_connection(self):
48+
"""
49+
Users can't get the default session without having a default connection set.
50+
"""
51+
with self.assertRaisesRegexp(connection.CQLEngineException, self.no_registered_connection_msg):
52+
connection.get_session(connection=None)
53+
54+
def test_get_cluster_fails_without_existing_connection(self):
55+
"""
56+
Users can't get the default cluster without having a default connection set.
57+
"""
58+
with self.assertRaisesRegexp(connection.CQLEngineException, self.no_registered_connection_msg):
59+
connection.get_cluster(connection=None)

‎tests/unit/cqlengine/test_udt.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Copyright 2013-2016 DataStax, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
try:
16+
import unittest2 as unittest
17+
except ImportError:
18+
import unittest # noqa
19+
20+
from cassandra.cqlengine import columns
21+
from cassandra.cqlengine.models import Model
22+
from cassandra.cqlengine.usertype import UserType
23+
24+
25+
class UDTTest(unittest.TestCase):
26+
27+
def test_initialization_without_existing_connection(self):
28+
"""
29+
Test that users can define models with UDTs without initializing
30+
connections.
31+
32+
Written to reproduce PYTHON-649.
33+
"""
34+
35+
class Value(UserType):
36+
t = columns.Text()
37+
38+
class DummyUDT(Model):
39+
__keyspace__ = 'ks'
40+
primary_key = columns.Integer(primary_key=True)
41+
value = columns.UserDefinedType(Value)

0 commit comments

Comments
 (0)