def register_connection(name, hosts=None, consistency=None, lazy_connect=False,
retry_connect=False, cluster_options=None, default=False,
session=None):
"""
Add a connection to the connection registry. ``hosts`` and ``session`` are
mutually exclusive, and ``consistency``, ``lazy_connect``,
``retry_connect``, and ``cluster_options`` only work with ``hosts``. Using
``hosts`` will create a new :class:`cassandra.cluster.Cluster` and
:class:`cassandra.cluster.Session`.
:param list hosts: list of hosts, (``contact_points`` for :class:`cassandra.cluster.Cluster`).
:param int consistency: The default :class:`~.ConsistencyLevel` for the
registered connection's new session. Default is the same as
:attr:`.Session.default_consistency_level`. For use with ``hosts`` only;
will fail when used with ``session``.
:param bool lazy_connect: True if should not connect until first use. For
use with ``hosts`` only; will fail when used with ``session``.
:param bool retry_connect: True if we should retry to connect even if there
was a connection failure initially. For use with ``hosts`` only; will
fail when used with ``session``.
:param dict cluster_options: A dict of options to be used as keyword
arguments to :class:`cassandra.cluster.Cluster`. For use with ``hosts``
only; will fail when used with ``session``.
:param bool default: If True, set the new connection as the cqlengine
default
:param Session session: A :class:`cassandra.cluster.Session` to be used in
the created connection.
"""
if name in _connections:
log.warning("Registering connection '{0}' when it already exists.".format(name))
if session is not None:
invalid_config_args = (hosts is not None or
consistency is not None or
lazy_connect is not False or
retry_connect is not False or
cluster_options is not None)
if invalid_config_args:
raise CQLEngineException(
"Session configuration arguments and 'session' argument are mutually exclusive"
)
conn = Connection.from_session(name, session=session)
conn.setup_session()
else: # use hosts argument
if consistency is None:
consistency = ConsistencyLevel.LOCAL_ONE
conn = Connection(
name, hosts=hosts,
consistency=consistency, lazy_connect=lazy_connect,
retry_connect=retry_connect, cluster_options=cluster_options
)
conn.setup()
_connections[name] = conn
if default:
set_default_connection(name)
return conn
def _inst_batch(self, batch):
assert self._timeout is connection.NOT_SET, 'Setting both timeout and batch is not supported'
if self._connection:
raise CQLEngineException(
"Cannot specify a connection on model in batch mode.")
self._batch = batch
return self
def set_session(s):
"""
Configures the default connection with a preexisting :class:`cassandra.cluster.Session`
Note: the mapper presently requires a Session :attr:`~.row_factory` set to ``dict_factory``.
This may be relaxed in the future
"""
try:
conn = get_connection()
except CQLEngineException:
# no default connection set; initalize one
register_connection('default', session=s, default=True)
conn = get_connection()
if conn.session:
log.warning("configuring new default connection for cqlengine when one was already set")
if s.row_factory is not dict_factory:
raise CQLEngineException("Failed to initialize: 'Session.row_factory' must be 'dict_factory'.")
conn.session = s
conn.cluster = s.cluster
# Set default keyspace from given session's keyspace
if conn.session.keyspace:
from cassandra.cqlengine import models
models.DEFAULT_KEYSPACE = conn.session.keyspace
conn.setup_session()
log.debug("cqlengine default connection initialized with %s", s)
async def _async_execute(self, statement):
connection = (self.instance._get_connection()
if self.instance else self.model._get_connection())
if self._batch:
if self._batch._connection:
if (not self._batch._connection_explicit and connection
and connection != self._batch._connection):
raise CQLEngineException(
"BatchQuery queries must be executed "
"on the same connection")
else:
# set the BatchQuery connection from the model
self._batch._connection = connection
return self._batch.add_query(statement)
else:
results = await _execute_statement(
self.model,
statement,
self._consistency,
self._timeout,
connection=connection,
)
if self._if_not_exists or self._if_exists or self._conditional:
check_applied(results)
return results
def setup(self):
"""Setup the connection"""
global cluster, session
if 'username' in self.cluster_options or 'password' in self.cluster_options:
raise CQLEngineException("Username & Password are now handled by using the native driver's auth_provider")
if self.lazy_connect:
return
self.cluster = Cluster(self.hosts, **self.cluster_options)
try:
self.session = self.cluster.connect()
log.debug(format_log_context("connection initialized with internally created session", connection=self.name))
except NoHostAvailable:
if self.retry_connect:
log.warning(format_log_context("connect failed, setting up for re-attempt on first use", connection=self.name))
self.lazy_connect = True
raise
if self.consistency is not None:
self.session.default_consistency_level = self.consistency
if DEFAULT_CONNECTION in _connections and _connections[DEFAULT_CONNECTION] == self:
cluster = _connections[DEFAULT_CONNECTION].cluster
session = _connections[DEFAULT_CONNECTION].session
self.setup_session()
def iff(self, *args, **kwargs):
"""Adds IF statements to queryset"""
if len([x for x in kwargs.values() if x is None]):
raise CQLEngineException("None values on iff are not allowed")
clone = copy.deepcopy(self)
for operator in args:
if not isinstance(operator, ConditionalClause):
raise QueryException('{0} is not a valid query operator'.format(operator))
clone._conditional.append(operator)
for arg, val in kwargs.items():
if isinstance(val, Token):
raise QueryException("Token() values are not valid in conditionals")
col_name, col_op = self._parse_filter_arg(arg)
try:
column = self.model._get_column(col_name)
except KeyError:
raise QueryException("Can't resolve column name: '{0}'".format(col_name))
if isinstance(val, BaseQueryFunction):
query_val = val
else:
query_val = column.to_database(val)
operator_class = BaseWhereOperator.get_operator(col_op or 'EQ')
operator = operator_class()
clone._conditional.append(WhereClause(column.db_field_name, operator, query_val))
return clone
def get_cluster():
handle_lazy_connect()
if not cluster:
raise CQLEngineException(
"%s.cluster is not configured. Call one of the setup or default functions first."
% __name__)
return cluster
def execute(query, params=None, consistency_level=None, timeout=NOT_SET):
handle_lazy_connect()
if not session:
raise CQLEngineException(
"It is required to setup() cqlengine before executing queries")
if isinstance(query, Statement):
pass
elif isinstance(query, BaseCQLStatement):
params = query.get_context()
query = str(query)
query = SimpleStatement(query, consistency_level=consistency_level)
elif isinstance(query, six.string_types):
query = SimpleStatement(query, consistency_level=consistency_level)
log.debug(query.query_string)
params = params or {}
result = session.execute(query, params, timeout=timeout)
return result
def __init__(self,
batch_type=None,
timestamp=None,
consistency=None,
execute_on_exception=False,
timeout=connection.NOT_SET):
"""
:param batch_type: (optional) One of batch type values available through BatchType enum
:type batch_type: str or None
:param timestamp: (optional) A datetime or timedelta object with desired timestamp to be applied
to the batch transaction.
:type timestamp: datetime or timedelta or None
:param consistency: (optional) One of consistency values ("ANY", "ONE", "QUORUM" etc)
:type consistency: The :class:`.ConsistencyLevel` to be used for the batch query, or None.
:param execute_on_exception: (Defaults to False) Indicates that when the BatchQuery instance is used
as a context manager the queries accumulated within the context must be executed despite
encountering an error within the context. By default, any exception raised from within
the context scope will cause the batched queries not to be executed.
:type execute_on_exception: bool
:param timeout: (optional) Timeout for the entire batch (in seconds), if not specified fallback
to default session timeout
:type timeout: float or None
"""
self.queries = []
self.batch_type = batch_type
if timestamp is not None and not isinstance(timestamp,
(datetime, timedelta)):
raise CQLEngineException(
'timestamp object must be an instance of datetime')
self.timestamp = timestamp
self._consistency = consistency
self._execute_on_exception = execute_on_exception
self._timeout = timeout
self._callbacks = []
def set_session(s):
"""
Configures the global mapper connection with a preexisting :class:`cassandra.cluster.Session`
Note: the mapper presently requires a Session :attr:`~.row_factory` set to ``dict_factory``.
This may be relaxed in the future
"""
global cluster, session
if session:
log.warning(
"configuring new connection for cqlengine when one was already set"
)
if s.row_factory is not dict_factory:
raise CQLEngineException(
"Failed to initialize: 'Session.row_factory' must be 'dict_factory'."
)
session = s
cluster = s.cluster
# Set default keyspace from given session's keyspace if not already set
from cassandra.cqlengine import models
if not models.DEFAULT_KEYSPACE and session.keyspace:
models.DEFAULT_KEYSPACE = session.keyspace
_register_known_types(cluster)
log.debug("cqlengine connection initialized with %s", s)
def test_connection_register_called_second_time(
self, connection_mock):
settings = self.cassandra_connection.settings_dict
CassandraConnection('default', **settings)
connection_mock.get_connection.side_effect = CQLEngineException()
self.assertFalse(connection_mock.register_connection.called)
def get_cluster(connection=None):
conn = get_connection(connection)
if not conn.cluster:
raise CQLEngineException(
"%s.cluster is not configured. Call one of the setup or default functions first."
% __name__)
return conn.cluster
def execute(query,
params=None,
consistency_level=None,
timeout=NOT_SET,
connection=None):
conn = get_connection(connection)
if not conn.session:
raise CQLEngineException(
"It is required to setup() cqlengine before executing queries")
if isinstance(query, SimpleStatement):
pass #
elif isinstance(query, BaseCQLStatement):
params = query.get_context()
query = SimpleStatement(str(query),
consistency_level=consistency_level,
fetch_size=query.fetch_size)
elif isinstance(query, six.string_types):
query = SimpleStatement(query, consistency_level=consistency_level)
log.debug(format_log_context(query.query_string, connection=connection))
result = conn.session.execute(query, params, timeout=timeout)
return result
def filter(self, *args, **kwargs):
"""
Adds WHERE arguments to the queryset, returning a new queryset
See :ref:`retrieving-objects-with-filters`
Returns a QuerySet filtered on the keyword arguments
"""
# add arguments to the where clause filters
if len([x for x in kwargs.values() if x is None]):
raise CQLEngineException("None values on filter are not allowed")
clone = copy.deepcopy(self)
for operator in args:
if not isinstance(operator, WhereClause):
raise QueryException('{0} is not a valid query operator'.format(operator))
clone._where.append(operator)
for arg, val in kwargs.items():
col_name, col_op = self._parse_filter_arg(arg)
quote_field = True
# resolve column and operator
try:
column = self.model._get_column(col_name)
except KeyError:
if col_name == 'pk__token':
if not isinstance(val, Token):
raise QueryException("Virtual column 'pk__token' may only be compared to Token() values")
column = columns._PartitionKeysToken(self.model)
quote_field = False
else:
raise QueryException("Can't resolve column name: '{0}'".format(col_name))
if isinstance(val, Token):
if col_name != 'pk__token':
raise QueryException("Token() values may only be compared to the 'pk__token' virtual column")
partition_columns = column.partition_columns
if len(partition_columns) != len(val.value):
raise QueryException(
'Token() received {0} arguments but model has {1} partition keys'.format(
len(val.value), len(partition_columns)))
val.set_columns(partition_columns)
# get query operator, or use equals if not supplied
operator_class = BaseWhereOperator.get_operator(col_op or 'EQ')
operator = operator_class()
if isinstance(operator, InOperator):
if not isinstance(val, (list, tuple)):
raise QueryException('IN queries must use a list/tuple value')
query_val = [column.to_database(v) for v in val]
elif isinstance(val, BaseQueryFunction):
query_val = val
else:
query_val = column.to_database(val)
clone._where.append(WhereClause(column.db_field_name, operator, query_val, quote_field=quote_field))
return clone
def _execute_query(self):
if self._batch:
raise CQLEngineException(
"Only inserts, updates, and deletes are available in batch mode"
)
if self._result_cache is None:
self._result_cache = list(self._execute(self._select_query()))
self._construct_result = self._get_result_constructor()
def setup(hosts,
default_keyspace,
consistency=ConsistencyLevel.ONE,
lazy_connect=False,
retry_connect=False,
**kwargs):
"""
Setup a the driver connection used by the mapper
:param list hosts: list of hosts, (``contact_points`` for :class:`cassandra.cluster.Cluster`)
:param str default_keyspace: The default keyspace to use
:param int consistency: The global default :class:`~.ConsistencyLevel`
:param bool lazy_connect: True if should not connect until first use
:param bool retry_connect: True if we should retry to connect even if there was a connection failure initially
:param \*\*kwargs: Pass-through keyword arguments for :class:`cassandra.cluster.Cluster`
"""
global cluster, session, default_consistency_level, lazy_connect_args
if 'username' in kwargs or 'password' in kwargs:
raise CQLEngineException(
"Username & Password are now handled by using the native driver's auth_provider"
)
if not default_keyspace:
raise UndefinedKeyspaceException()
from cassandra.cqlengine import models
models.DEFAULT_KEYSPACE = default_keyspace
default_consistency_level = consistency
if lazy_connect:
kwargs['default_keyspace'] = default_keyspace
kwargs['consistency'] = consistency
kwargs['lazy_connect'] = False
kwargs['retry_connect'] = retry_connect
lazy_connect_args = (hosts, kwargs)
return
cluster = Cluster(hosts, **kwargs)
try:
session = cluster.connect()
log.debug(
"cqlengine connection initialized with internally created session")
except NoHostAvailable:
if retry_connect:
log.warning(
"connect failed, setting up for re-attempt on first use")
kwargs['default_keyspace'] = default_keyspace
kwargs['consistency'] = consistency
kwargs['lazy_connect'] = False
kwargs['retry_connect'] = retry_connect
lazy_connect_args = (hosts, kwargs)
raise
session.row_factory = dict_factory
_register_known_types(cluster)
def _get_partition_keys(self):
try:
table_meta = get_cluster(self._get_connection()).metadata.keyspaces[self.keyspace].tables[self.name]
self.__partition_keys = OrderedDict(
(pk.name, Column(primary_key=True, partition_key=True, db_field=pk.name)) for pk in
table_meta.partition_key)
except Exception as e:
raise CQLEngineException("Failed inspecting partition keys for {0}."
"Ensure cqlengine is connected before attempting this with NamedTable.".format(
self.column_family_name()))
def set_default_connection(name):
global cluster, session
if name not in _connections:
raise CQLEngineException("Connection '{0}' doesn't exist.".format(name))
log.debug("Connection '{0}' has been set as default.".format(name))
_connections[DEFAULT_CONNECTION] = _connections[name]
cluster = _connections[name].cluster
session = _connections[name].session
def batch(self, batch_obj):
"""
Set a batch object to run the query on.
Note: running a select query with a batch object will raise an exception
"""
if batch_obj is not None and not isinstance(batch_obj, BatchQuery):
raise CQLEngineException('batch_obj must be a BatchQuery instance or None')
clone = copy.deepcopy(self)
clone._batch = batch_obj
return clone
def delete(self):
""" Deletes one instance """
if self.instance is None:
raise CQLEngineException("DML Query instance attribute is None")
ds = DeleteStatement(self.column_family_name, timestamp=self._timestamp, conditionals=self._conditional, if_exists=self._if_exists)
for name, col in self.model._primary_keys.items():
val = getattr(self.instance, name)
if val is None and not col.partition_key:
continue
ds.add_where(col, EqualsOperator(), val)
self._execute(ds)
def _execute_query(self):
if self._batch:
raise CQLEngineException("Only inserts, updates, and deletes are available in batch mode")
if self._result_cache is None:
self._result_generator = (i for i in self._execute(self._select_query()))
self._result_cache = []
self._construct_result = self._maybe_inject_deferred(self._get_result_constructor())
# "DISTINCT COUNT()" is not supported in C* < 2.2, so we need to materialize all results to get
# len() and count() working with DISTINCT queries
if self._materialize_results or self._distinct_fields:
self._fill_result_cache()
def get_connection(name=None):
if not name:
name = DEFAULT_CONNECTION
if name not in _connections:
raise CQLEngineException("Connection name '{0}' doesn't exist in the registry.".format(name))
conn = _connections[name]
conn.handle_lazy_connect()
return conn
async def async_save(self):
"""
Creates / updates a row.
This is a blind insert call.
All validation and cleaning needs to happen
prior to calling this.
"""
if self.instance is None:
raise CQLEngineException("DML Query intance attribute is None")
assert type(self.instance) == self.model
nulled_fields = set()
if self.instance._has_counter or self.instance._can_update():
if self.instance._has_counter:
warn("'create' and 'save' actions on Counters are deprecated. "
"A future version will disallow this. Use the 'update' "
"mechanism instead.")
return await self.async_update()
else:
insert = InsertStatement(
self.column_family_name,
ttl=self._ttl,
timestamp=self._timestamp,
if_not_exists=self._if_not_exists,
)
static_save_only = (False if len(self.instance._clustering_keys)
== 0 else True)
for name, col in self.instance._clustering_keys.items():
static_save_only = static_save_only and col._val_is_null(
getattr(self.instance, name, None))
for name, col in self.instance._columns.items():
if (static_save_only and not col.static
and not col.partition_key):
continue
val = getattr(self.instance, name, None)
if col._val_is_null(val):
if self.instance._values[name].changed:
nulled_fields.add(col.db_field_name)
continue
if col.has_default and not self.instance._values[name].changed:
# Ensure default columns included in a save()
# are marked as explicit, to get them *persisted* properly
self.instance._values[name].explicit = True
insert.add_assignment(col, getattr(self.instance, name, None))
# skip query execution if it's empty
# caused by pointless update queries
if not insert.is_empty:
await self._async_execute(insert)
# delete any nulled columns
if not static_save_only:
self._delete_null_columns()
def column_family_name(cls, include_keyspace=True):
"""
Returns the column family name if it's been defined
otherwise, it creates it from the module and class name
"""
cf_name = protect_name(cls._raw_column_family_name())
if include_keyspace:
keyspace = cls._get_keyspace()
if not keyspace:
raise CQLEngineException("Model keyspace is not set and no default is available. Set model keyspace or setup connection before attempting to generate a query.")
return '{0}.{1}'.format(protect_name(keyspace), cf_name)
return cf_name
def test_register_connection_called_first_time_with_proper_options(
self, connection_mock):
settings = self.cassandra_connection.settings_dict
connection = CassandraConnection('default', **settings)
connection_mock.get_connection.side_effect = CQLEngineException()
connection.register()
connection_mock.register_connection.assert_called_once_with(
'default',
hosts=connection.hosts,
default=mock.ANY,
cluster_options=connection.cluster_options,
**connection.connection_options)
def _validate_pk(model, table_meta):
model_partition = [c.db_field_name for c in model._partition_keys.values()]
meta_partition = [c.name for c in table_meta.partition_key]
model_clustering = [c.db_field_name for c in model._clustering_keys.values()]
meta_clustering = [c.name for c in table_meta.clustering_key]
if model_partition != meta_partition or model_clustering != meta_clustering:
def _pk_string(partition, clustering):
return "PRIMARY KEY (({0}){1})".format(', '.join(partition), ', ' + ', '.join(clustering) if clustering else '')
raise CQLEngineException("Model {0} PRIMARY KEY composition does not match existing table {1}. "
"Model: {2}; Table: {3}. "
"Update model or drop the table.".format(model, model.column_family_name(),
_pk_string(model_partition, model_clustering),
_pk_string(meta_partition, meta_clustering)))
def count(self):
"""
Returns the number of rows matched by this query
"""
if self._batch:
raise CQLEngineException("Only inserts, updates, and deletes are available in batch mode")
if self._result_cache is None:
query = self._select_query()
query.count = True
result = self._execute(query)
return result[0]['count']
else:
return len(self._result_cache)
def update(self):
"""
updates a row.
This is a blind update call.
All validation and cleaning needs to happen
prior to calling this.
"""
if self.instance is None:
raise CQLEngineException("DML Query intance attribute is None")
assert type(self.instance) == self.model
null_clustering_key = False if len(self.instance._clustering_keys) == 0 else True
static_changed_only = True
statement = UpdateStatement(self.column_family_name, ttl=self._ttl, timestamp=self._timestamp,
conditionals=self._conditional, if_exists=self._if_exists)
for name, col in self.instance._clustering_keys.items():
null_clustering_key = null_clustering_key and col._val_is_null(getattr(self.instance, name, None))
updated_columns = set()
# get defined fields and their column names
for name, col in self.model._columns.items():
# if clustering key is null, don't include non static columns
if null_clustering_key and not col.static and not col.partition_key:
continue
if not col.is_primary_key:
val = getattr(self.instance, name, None)
val_mgr = self.instance._values[name]
if val is None:
continue
if not val_mgr.changed and not isinstance(col, columns.Counter):
continue
static_changed_only = static_changed_only and col.static
statement.add_update(col, val, previous=val_mgr.previous_value)
updated_columns.add(col.db_field_name)
if statement.assignments:
for name, col in self.model._primary_keys.items():
# only include clustering key if clustering key is not null, and non static columns are changed to avoid cql error
if (null_clustering_key or static_changed_only) and (not col.partition_key):
continue
statement.add_where(col, EqualsOperator(), getattr(self.instance, name))
self._execute(statement)
if not null_clustering_key:
# remove conditions on fields that have been updated
delete_conditionals = [condition for condition in self._conditional
if condition.field not in updated_columns] if self._conditional else None
self._delete_null_columns(delete_conditionals)
def __init__(self, model, keyspace=None):
"""
:param model: A model. This should be a class type, not an instance.
:param keyspace: (optional) A keyspace name
"""
from cassandra.cqlengine import models
if not issubclass(model, models.Model):
raise CQLEngineException("Models must be derived from base Model.")
ks = keyspace if keyspace else model.__keyspace__
new_type = type(model.__name__, (model,), {'__keyspace__': ks})
self.model = new_type
def __get__(self, obj, model):
""" :rtype: ModelQuerySet """
if model.__abstract__:
raise CQLEngineException('cannot execute queries against abstract models')
queryset = model.__queryset__(model)
# if this is a concrete polymorphic model, and the discriminator
# key is an indexed column, add a filter clause to only return
# logical rows of the proper type
if model._is_polymorphic and not model._is_polymorphic_base:
name, column = model._discriminator_column_name, model._discriminator_column
if column.partition_key or column.index:
# look for existing poly types
return queryset.filter(**{name: model.__discriminator_value__})
return queryset