def _get_sb_fields(self, fields):
"""Return sqlbuilder columns for a SELECT query based on
passed field names or * if no fields are passed
Prefers native columns where available
See doctest_PJContainer_get_sb_fields"""
if not fields:
res = sb.Field(self._pj_table, '*')
else:
datafld = sb.Field(self._pj_table, self._pj_data_column)
res = []
# prefer sql columns over json fields
for name in fields:
if name in self._pj_column_fields:
# SQL names are case-insensitive, so when we return the
# result, we want proper capitalization.
res.append(sb.ColumnAS(
sb.Field(self._pj_table, name.lower()), name))
else:
if '.' not in name:
accessor = sb.JSON_GETITEM(datafld, name)
else:
accessor = sb.JSON_PATH(datafld, name.split("."))
res.append(sb.ColumnAS(accessor, name.replace('.', '_')))
return res
def _pj_get_pj_mapping_key_field(self):
"""Return an sqlbuilder field or JSON field getter to be used
to get the value of the _pj_get_pj_mapping field"""
if self._pj_mapping_key in self._pj_column_fields:
# if it's a native column, no need to dig in JSONB
fld = sb.Field(self._pj_table, self._pj_mapping_key)
else:
datafld = sb.Field(self._pj_table, self._pj_data_column)
fld = sb.JGET(datafld, self._pj_mapping_key)
return fld
def __nonzero__(self):
where = self._pj_get_list_filter() or True
select = sb.Select(sb.func.COUNT(
sb.Field(self._pj_table, self._pj_id_column)), where=where)
with self._pj_jar.getCursor() as cur:
cur.execute(select, flush_hint=[self._pj_table])
return cur.fetchone()[0] > 0
def _get_doc(self, database, table, id):
tbl = sb.Table(table)
with self.getCursor() as cur:
cur.execute(sb.Select(sb.Field(table, '*'), tbl.id == id),
beacon="%s:%s:%s" % (database, table, id),
flush_hint=[table])
res = cur.fetchone()
return res['data'] if res is not None else None
def _get_doc_py_type(self, database, table, id):
tbl = sb.Table(table)
with self.getCursor() as cur:
datafld = sb.Field(table, 'data')
cur.execute(sb.Select(
sb.JGET(datafld, interfaces.PY_TYPE_ATTR_NAME), tbl.id == id),
beacon="%s:%s:%s" % (database, table, id))
res = cur.fetchone()
return res[0] if res is not None else None
def __getitem__(self, key):
with self._jar.getCursor(False) as cur:
tbl = sb.Table(self.table)
cur.execute(
sb.Select(sb.Field(self.table, 'dbref'), tbl.name == key))
if not cur.rowcount:
raise KeyError(key)
db, tbl, id = cur.fetchone()['dbref']
dbref = serialize.DBRef(tbl, id, db)
return self._jar.load(dbref)
def select(conn, query, print_sql=False, **kwargs):
try:
with conn.cursor() as cur:
sql = sb.sqlrepr(
sb.Select(sb.Field("ser", "data"), where=query, **kwargs),
'postgres')
if print_sql:
print('SQL> ', sql)
cur.execute(sql)
for e in cur.fetchall():
pprint(e[0])
finally:
conn.rollback()
def select(conn, query, print_sql=False):
try:
with conn.cursor() as cur:
converter = mquery.Converter("mq", "data")
sql = sb.sqlrepr(
sb.Select(sb.Field("mq", "data"),
where=converter.convert(query)), 'postgres')
if print_sql:
print('SQL> ', sql)
cur.execute(sql)
for e in cur.fetchall():
pprint.pprint(e[0])
finally:
conn.rollback()
def count(self, qry=None):
if isinstance(qry, dict):
qry = self.convert_mongo_query(qry)
where = self._combine_filters(self._pj_get_list_filter(), qry)
count = sb.func.COUNT(sb.Field(self._pj_table, self._pj_id_column))
if where is None:
select = sb.Select(count)
else:
select = sb.Select(count, where=where)
with self._pj_jar.getCursor() as cur:
cur.execute(select, flush_hint=[self._pj_table])
return cur.fetchone()[0]
def getField(self, field, key, json=False):
if isinstance(field, sb.SQLConstant):
# hack for $elemMatch
accessor = field
elif key == '_id':
accessor = sb.Field(self.table, 'id')
elif json:
if '.' not in key:
accessor = sb.JSON_GETITEM(field, key)
else:
accessor = sb.JSON_PATH(field, key.split("."))
else:
if '.' not in key:
accessor = sb.JSON_GETITEM_TEXT(field, key)
else:
accessor = sb.JSON_PATH_TEXT(field, key.split("."))
return accessor
def convert(self, query):
clauses = []
doc = sb.Field(self.table, self.field)
for key, value in sorted(query.items()):
accessor = self.getField(doc, key, json=True)
# some values, esp. datetime must go through PJ serialize
pjvalue = serialize.ObjectWriter(None).get_state(value)
jvalue = json.dumps(pjvalue, sort_keys=True)
if key == '_id':
jvalue = value
if key in ('$and', '$or', '$nor', '$startswith'):
if not isinstance(value, (list, tuple)):
raise ValueError("Argument must be a list: %r" % value)
oper = {
'$and': sb.AND,
'$or': sb.OR,
'$nor': lambda *items: sb.NOT(sb.OR(*items)),
'$startswith': sb.STARTSWITH # special case for a $regex
}[key]
clauses.append(oper(*(self.convert(expr) for expr in value)))
elif isinstance(value, dict):
for operator, operand in sorted(value.items()):
clauses.append(
self.operator_expr(operator, doc, key, operand))
else:
# Scalar -- equality or array membership
if self.simplified or key == '_id':
# Let's ignore the membership case for test clarity
clauses.append(accessor == jvalue)
else:
options = [
accessor == jvalue,
sb.JSONB_SUBSET(sb.JSONB(json.dumps([pjvalue])),
accessor)
]
if value is None:
options.append(sb.ISNULL(accessor))
clauses.append(sb.OR(*options))
return sb.AND(*clauses)
def _pj_get_resolve_filter(self):
"""return a filter that selects the rows of the current container"""
queries = []
# Make sure that we only look through objects that have the mapping
# key. Objects not having the mapping key cannot be part of the
# table.
datafld = sb.Field(self._pj_table, self._pj_data_column)
if self._pj_mapping_key is not None:
if self._pj_mapping_key not in self._pj_column_fields:
# if `_pj_mapping_key` is a JSONB field make sure that it
# exists, JGET returns jsonb which is never NULL
# if `_pj_mapping_key` is a native column, you need to make
# sure that it exists
queries.append(
sb.JSONB_CONTAINS(datafld, self._pj_mapping_key))
# We also make want to make sure we separate the items properly by the
# container.
if self._pj_parent_key is not None:
pv = self._pj_jar._writer.get_state(self._pj_get_parent_key_value())
queries.append(sb.JGET(datafld, self._pj_parent_key) == pv)
return sb.AND(*queries)
def raw_find_one(self, qry=None, id=None):
if qry is None and id is None:
raise ValueError(
'Missing parameter, at least qry or id must be specified.')
if isinstance(qry, dict):
qry = self.convert_mongo_query(qry)
if id is not None:
qry = self._combine_filters(
self._pj_get_resolve_filter(), qry, self._pj_get_id_filter(id))
else:
qry = self._combine_filters(
self._pj_get_list_filter(), qry)
with self._pj_jar.getCursor() as cur:
cur.execute(sb.Select(sb.Field(self._pj_table, '*'), qry, limit=2),
flush_hint=[self._pj_table])
if cur.rowcount == 0:
return None
if cur.rowcount > 1:
raise ValueError('Multiple results returned.')
return cur.fetchone()
def keys(self):
with self._jar.getCursor(False) as cur:
cur.execute(sb.Select(sb.Field(self.table, 'name')))
return [doc['name'] for doc in cur.fetchall()]
