def test_insert_returning_all_fields_and_arithmetics(self):
query = PostgreSQLQuery.into(self.table_abc).insert(1).returning(
self.table_abc.star,
self.table_abc.f1 + self.table_abc.f2
)
self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING *,f1+f2', str(query))
def test_postgres_query_uses_double_quote_chars(self):
q = PostgreSQLQuery.from_(self.table_abc).select(
self.table_abc.foo
).groupby(
self.table_abc.foo
).having(
self.table_abc.buz == 'fiz'
)
self.assertEqual("SELECT \"foo\" FROM \"abc\" GROUP BY \"foo\" HAVING \"buz\"='fiz'", str(q))
def test_insert_all_columns_multi_rows_chained_returning_star_and_id(self):
query = PostgreSQLQuery.into(
self.table_abc
).insert(1, 'a', True).insert(2, 'b', False).returning(
self.table_abc.name,
self.table_abc.star,
self.table_abc.id,
)
self.assertEqual('INSERT INTO "abc" VALUES (1,\'a\',true),(2,\'b\',false) RETURNING *', str(query))
def test_on_conflict_where_conflict_fields_do_nothing(self):
qs = (
PostgreSQLQuery.into(self.table_abc)
.insert(1, "m")
.on_conflict('id')
.where(self.table_abc.abc.eq(0))
.where(self.table_abc.cde.eq(0))
.do_nothing()
)
self.assertEqual(
'''INSERT INTO "abc" VALUES (1,'m') ON CONFLICT ("id") WHERE "abc"=0 AND "cde"=0 DO NOTHING''', str(qs))
def _rem_mod(itgs: LazyItgs, user_id: int, commit: bool = False):
moderators = Table('moderators')
itgs.write_cursor.execute(
Query.from_(moderators)
.delete()
.where(moderators.user_id == Parameter('%s'))
.get_sql(),
(user_id,)
)
if commit:
itgs.write_conn.commit()
def _add_mod(itgs: LazyItgs, user_id: int, commit: bool = False):
moderators = Table('moderators')
itgs.write_cursor.execute(
Query.into(moderators)
.columns(moderators.user_id)
.insert(Parameter('%s'))
.get_sql(),
(user_id,)
)
if commit:
itgs.write_conn.commit()
def delete_game(self, chat_id: int) -> None:
if not self.use_db:
return
games = Table('games')
logging.info(f"Removing game {chat_id} from DB")
delete_game_query = PostgreSQLQuery.from_(games) \
.delete() \
.where(games.chat_id == chat_id)
self._cur.execute(delete_game_query.get_sql())
self._conn.commit()
def fetch_course_subject_prefixes(conn: RealDictConnection) -> List[str]:
cursor = conn.cursor()
q: QueryBuilder = (
Query.from_(course_sections_t)
.select(course_sections_t.course_subject_prefix)
.groupby(course_sections_t.course_subject_prefix)
.orderby(course_sections_t.course_subject_prefix)
)
cursor.execute(q.get_sql())
return list(map(lambda record: record["course_subject_prefix"], cursor.fetchall()))
def test_insert_on_conflict_do_nothing_mixed_fields(self):
query = (
PostgreSQLQuery.into(self.table_abc)
.insert(1)
.on_conflict("id", self.table_abc.sub_id)
.do_nothing()
)
self.assertEqual(
'INSERT INTO "abc" VALUES (1) ON CONFLICT ("id", "sub_id") DO NOTHING', str(query)
)
def test_insert_all_columns_multi_rows_chained_returning_star_and_id(self):
query = PostgreSQLQuery.into(self.table_abc).insert(
1, 'a', True).insert(2, 'b', False).returning(
self.table_abc.name,
self.table_abc.star,
self.table_abc.id,
)
self.assertEqual(
'INSERT INTO "abc" VALUES (1,\'a\',true),(2,\'b\',false) RETURNING *',
str(query))
def test_insert_on_fieldless_conflict_do_nothing(self):
query = (
PostgreSQLQuery.into(self.table_abc)
.insert(1)
.on_conflict(None)
.do_nothing()
)
self.assertEqual(
'INSERT INTO "abc" VALUES (1) ON CONFLICT DO NOTHING', str(query)
)
def get_loansbot_user_id(itgs: 'LazyItgs') -> int:
"""This function returns the id of the loansbot user. If they do not exist
they are created. This is useful since thats the "author" of automated
permission changes.
Arguments:
- `itgs (LazyItgs)`: The integrations for connecting to networked
services.
Returns:
- `loansbot_user_id (int)`: The id of the loansbot user.
"""
users = Table('users')
unm = 'loansbot'
(user_id, ) = query_helper.find_or_create_or_find(
itgs, (Query.from_(users).select(
users.id).where(users.username == Parameter('%s')).get_sql(),
(unm, )), (Query.into(users).columns(users.username).insert(
Parameter('%s')).returning(users.id).get_sql(), (unm, )))
return user_id
def get_permissions_from_header(itgs, authorization, permissions):
"""A convenience method to get if authorization was provided, if it was
valid, and which (if any) of the specified permissions they have. This is
a good balance of convenience and versatility.
Example responses:
Omitted authorization header
(None, False, [])
Invalid authorization header
(None, True, [])
Valid authorization header for user 3, no permissions from list
(3, True, [])
Valid authorization header for user 7, with some permissions
(7, True, ['permission1', 'permission2'])
@param itgs The lazy integrations to use
@param authorization The authorization header provided
@param permissions The list of interesting permissions for this endpoint;
this will return the subset of these permissions which the user
actually has.
@return (int, bool, list) If they authorized successfully then the user id
otherwise None, if they provided an authorization header, and what
permissions from the given list of permissions they have.
"""
if isinstance(permissions, str):
permissions = [permissions]
authtoken = get_authtoken_from_header(authorization)
if authtoken is None:
return (None, False, [])
info = get_auth_info_from_token_auth(
itgs, models.TokenAuthentication(token=authtoken)
)
if info is None:
return (None, True, [])
auth_id, user_id = info[:2]
if not permissions:
return (user_id, True, [])
perms = Table('permissions')
authtoken_perms = Table('authtoken_permissions')
itgs.read_cursor.execute(
Query.from_(authtoken_perms).select(perms.name)
.join(perms).on(perms.id == authtoken_perms.permission_id)
.where(perms.name.isin([Parameter('%s') for _ in permissions]))
.where(authtoken_perms.authtoken_id == Parameter('%s'))
.get_sql(),
(*permissions, auth_id)
)
perms_found = itgs.read_cursor.fetchall()
return (user_id, True, [i[0] for i in perms_found])
def test_on_conflict_do_update_where(self):
qs = (
PostgreSQLQuery.into(self.table_abc)
.insert(1, "m")
.on_conflict("id")
.do_update('abc', 1)
.where(self.table_abc.abc.eq(1))
)
self.assertEqual(
'INSERT INTO "abc" VALUES (1,\'m\') ON CONFLICT ("id") DO UPDATE SET "abc"=1 WHERE "abc"."abc"=1', str(qs)
)
async def update_chat(self, request, context, metadata):
query = PostgreSQLQuery.update(self.chats_table)
for field in request.DESCRIPTOR.fields:
if field.containing_oneof and request.HasField(field.name):
field_value = getattr(request, field.name)
query = query.set(field.name, field_value)
query = query.where(self.chats_table.chat_id ==
request.chat_id).returning('*').get_sql()
async with self.pool_holder.pool.acquire() as session:
result = await session.execute(query)
chat = await result.fetchone()
return self.enrich_chat(ChatPb(**chat))
def write_dict_to_db(connection, response, pg_table):
# pprint(response)
# response2 = add_exported_already_field(response, pg_table)
# response3 = add_date_exported_field(response2, pg_table)
cursor = connection.cursor()
pg_query2 = PostgreSQLQuery.into(pg_table).insert(*response.values())
pg_query = pg_query2.on_conflict().do_nothing()
# pprint(pg_query2)
cursor.execute(pg_query.get_sql())
connection.commit()
cursor.close()
return cursor.rowcount
def load_db_table(conn, table_name) -> List[Dict]:
channels = Table(table_name)
q = PostgreSQLQuery.from_(channels).select(channels.star)
cur = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
cur.execute(q.get_sql())
ans = cur.fetchall()
ans1 = []
for row in ans:
ans1.append(dict(row))
return ans1
def test_insert_on_conflict_do_update_multiple_mixed_fields(self):
query = (
PostgreSQLQuery.into(self.table_abc)
.insert(1, "m")
.on_conflict("id", self.table_abc.sub_id)
.do_update("name", "m")
)
self.assertEqual(
"INSERT INTO \"abc\" VALUES (1,'m') ON CONFLICT (\"id\", \"sub_id\") DO UPDATE SET \"name\"='m'",
str(query),
)
def test_insert_on_conflict_do_update_field(self):
query = (
PostgreSQLQuery.into(self.table_abc)
.insert(1, "m")
.on_conflict(self.table_abc.id)
.do_update(self.table_abc.name, "m")
)
self.assertEqual(
"INSERT INTO \"abc\" VALUES (1,'m') ON CONFLICT (\"id\") DO UPDATE SET \"name\"='m'",
str(query),
)
def test_insert_all_columns_multi_rows_chained_returning_star_and_id(self):
query = (PostgreSQLQuery.into(self.table_abc).insert(
1, "a", True).insert(2, "b", False).returning(
self.table_abc.name,
self.table_abc.star,
self.table_abc.id,
))
self.assertEqual(
"INSERT INTO \"abc\" VALUES (1,'a',true),(2,'b',false) RETURNING *",
str(query),
)
def window(self,
window_fn,
term=None,
over=None,
order_by=None,
columns=None,
**kwargs):
self.current_query = Query.from_(self.current_dataset)
tmp_query = ''
window_ = getattr(an, window_fn)
argumentless_analytics = ['DenseRank', 'Rank', 'RowNumber']
positional_arg_analytics = ['FirstValue', 'LastValue']
term_analytics = [
'NTile', 'Median', 'Avg', 'StdDev', 'StdDevPop', 'StdDevSamp',
'Variance', 'VarPop', 'VarSamp', 'Count', 'Sum', 'Max', 'Min'
]
if window_fn in argumentless_analytics:
tmp_query = window_()
elif window_fn in positional_arg_analytics: # Only uses first positional arg. Cannot use `term=`
tmp_query = window_(getattr(self.current_dataset, term))
elif window_fn in term_analytics:
try:
tmp_query = window_(term=getattr(self.current_dataset, term),
**kwargs)
except TypeError: # Term is an integer, like with NTile
tmp_query = window_(term=term, **kwargs)
else:
raise Exception(
'window_fn was not found in predefined window functions.')
if over:
for over_elem in over:
tmp_query = tmp_query.over(
getattr(self.current_dataset, over_elem))
if order_by:
for order in order_by:
tmp_query = tmp_query.orderby(
getattr(self.current_dataset, order))
if columns:
self.current_query = self.current_query.select(
*[
getattr(self.current_dataset, col_elem)
for col_elem in columns
], tmp_query)
else:
self.current_query = self.current_query.select(
self.current_dataset.star, tmp_query)
self.current_dataset = self.current_query
return self
async def createOrder(self, order: Order.CreateInfo) -> int:
if len(order.products) == 0: raise Exception('No products')
orderSql = Query.into(orders).columns(
orders.address, orders.name, orders.surname, orders.user_id,
orders.status, orders.shipping_price).insert(
order.address, order.name, order.surname, order.userId,
Cast('NEW', 'ORDER_STATUS'),
order.shippingPrice).returning(orders.id)
productsSql = Query.into(order_product_link).columns(
order_product_link.order_id, order_product_link.product_id,
order_product_link.amount)
async with databaseClient.transaction() as t:
orderId = (await t.query(orderSql.get_sql()))[0][0]
for prodId, amount in order.products.items():
productsSql = productsSql.insert(orderId, prodId, amount)
await t.execute(productsSql.get_sql())
return orderId
def test_update_returning_from_different_tables(self):
table_bcd = Table('bcd')
q = (
PostgreSQLQuery.update(self.table_abc)
.from_(table_bcd)
.set(self.table_abc.lname, table_bcd.long_name)
.returning(self.table_abc.id, table_bcd.fname)
)
self.assertEqual(
'UPDATE "abc" SET "lname"="bcd"."long_name" FROM "bcd" RETURNING "abc"."id","bcd"."fname"', str(q)
)
async def _get_id(zoid):
storage = get_storage()
if storage is None:
return
async with storage.pool.acquire() as conn:
result = await conn.fetch(
str(
Query.from_(objects_table).select(
"id",
"parent_id").where(objects_table.zoid == _safe_uid(zoid))))
if len(result) > 0:
return {"id": result[0]["id"], "parent": result[0]["parent_id"]}
def test_on_conflict_where_conflict_fields_do_update(self):
qs = (
PostgreSQLQuery.into(self.table_abc)
.insert(1, "m")
.on_conflict('id')
.where(self.table_abc.abc.eq(0))
.where(self.table_abc.cde.eq(0))
.do_update('field', 'val')
)
self.assertEqual(
'''INSERT INTO "abc" VALUES (1,'m') ON CONFLICT ("id") WHERE "abc"=0 AND "cde"=0 '''
'''DO UPDATE SET "field"='val\'''', str(qs))
def bonus_filters(query, add_param, endpoint_users, **kwargs):
endpoints = Table('endpoints')
return (query.where(
Exists(
Query.from_(endpoints).where(
endpoints.id == endpoint_users.endpoint_id).
where(endpoints.sunsets_on > Now() - Interval(days=27)).where(
endpoints.sunsets_on < Now()).where(
DatePart('day', endpoints.sunsets_on -
Now()) < (30 - Floor(
DatePart(
'day', endpoint_users.created_at -
endpoints.sunsets_on) / 3) * 3)))))
def generate_insert_sql(self, document_id: int, value: int, voter_id: int):
query = PostgreSQLQuery.into(self.votes_table).columns(
'document_id',
'value',
'voter_id',
).insert(
document_id,
value,
voter_id,
)
query = query.on_conflict('document_id',
'voter_id').do_update('value', value)
return query.get_sql()
def clear_temporary_bans(itgs: LazyItgs, username: str, subreddit: str) -> None:
temp_bans = Table('temporary_bans')
users = Table('users')
itgs.write_cursor.execute(
Query.from_(temp_bans).delete()
.where(Exists(
Query.from_(users).select(1)
.where(users.id == temp_bans.user_id)
.where(users.username == Parameter('%s'))
))
.where(temp_bans.subreddit == Parameter('%s'))
.get_sql(),
(username, subreddit)
)
itgs.write_conn.commit()
itgs.logger.print(
Level.DEBUG,
'Deleted temporary bans for /u/{} on /r/{} (overriden by another '
+ 'mod action)',
username, subreddit,
)
def load_unused_users(conn, table_name) -> List[Dict]:
channels = Table(table_name)
q = PostgreSQLQuery.from_(channels).select(channels.username, channels.has_been_used).where(
channels.has_been_used == False)
# pprint(q)
cur = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
cur.execute(q.get_sql())
ans = cur.fetchall()
ans1 = []
for row in ans:
ans1.append(dict(row))
return ans1
def get_column_definitions(self, schema, table, connection=None):
columns = Table("columns", schema="INFORMATION_SCHEMA")
columns_query = (
PostgreSQLQuery.from_(columns, immutable=False)
.select(columns.column_name, columns.data_type)
.where(columns.table_schema == schema)
.where(columns.field("table_name") == table)
.distinct()
.orderby(columns.column_name)
)
return self.fetch(str(columns_query), connection=connection)
def test_where_and_on_conflict_where(self):
table_bcd = Table('bcd')
qs = (PostgreSQLQuery.into(
self.table_abc).select(table_bcd.abc).from_(table_bcd).where(
table_bcd.abc.eq('1')).on_conflict('id').where(
self.table_abc.abc.eq(0)).where(
self.table_abc.cde.eq(0)).do_update('field', 'val'))
self.assertEqual(
'INSERT INTO "abc" SELECT "abc" FROM "bcd" WHERE "abc"=\'1\' '
'ON CONFLICT ("id") WHERE "abc"=0 AND "cde"=0 DO UPDATE SET "field"=\'val\'',
str(qs))
def get_max_rv_point_list(self, point_table: str) -> int:
rv = 0
try:
p = Table(point_table)
q = (Q.from_(p).select(Max(p.rv)))
sql = q.get_sql()
self._logger.debug(f'SQL: {sql}')
rv = engine_mart.scalar(sql)
if rv is None:
rv = 0
except Exception as e:
self._logger.error(e)
return rv
def find_or_create_money(itgs: LazyItgs, money: Money,
amount_usd_cents: int) -> int:
"""Find or create a row in the moneys table that matches the given money
object.
Arguments:
- `itgs (LazyItgs)`: The integrations to use to connect to networked
services.
- `money (Money)`: The money object to create a corresponding row for.
- `amount_usd_cents (int)`: The money amount in USD cents at the current
conversion rate.
"""
currencies = Table('currencies')
moneys = Table('moneys')
(currency_id, ) = query_helper.find_or_create_or_find(
itgs, (Query.from_(currencies).select(
currencies.id).where(currencies.code == Parameter('%s')).get_sql(),
(money.currency, )),
(Query.into(currencies).columns(
currencies.code, currencies.symbol, currencies.symbol_on_left,
currencies.exponent).insert(*(Parameter('%s')
for _ in range(4))).returning(
currencies.id).get_sql(),
(money.currency, money.symbol or f' {money.currency}',
money.symbol_on_left if money.symbol is not None else False,
money.exp or 2)))
(money_id, ) = query_helper.find_or_create_or_find(
itgs, (Query.from_(moneys).select(
moneys.id).where(moneys.currency_id == Parameter('%s')).where(
moneys.amount == Parameter('%s')).where(
moneys.amount_usd_cents == Parameter('%s')).get_sql(),
(currency_id, money.minor, amount_usd_cents)),
(Query.into(moneys).columns(
moneys.currency_id, moneys.amount, moneys.amount_usd_cents).insert(
*(Parameter('%s')
for _ in range(3))).returning(moneys.id).get_sql(),
(currency_id, money.minor, amount_usd_cents)))
return money_id
def test_postgresql_query_uses_double_quote_chars(self):
q = PostgreSQLQuery.from_('abc').select('foo', 'bar')
self.assertEqual('SELECT "foo","bar" FROM "abc"', str(q))
def test_update_returning(self):
q = PostgreSQLQuery.update(self.table_abc).where(
self.table_abc.foo == 0
).set('foo', 'bar').returning(self.table_abc.id)
self.assertEqual('UPDATE "abc" SET "foo"=\'bar\' WHERE "foo"=0 RETURNING id', str(q))
def test_insert_on_conflict_do_update_field_str(self):
query = PostgreSQLQuery.into(self.table_abc).insert(1, "m").on_conflict(
'id').do_update('name', "m")
self.assertEqual('INSERT INTO "abc" VALUES (1,\'m\') ON CONFLICT (id) DO UPDATE SET name=\'m\'', str(query))
def test_insert_on_conflict_no_handler(self):
with self.assertRaises(QueryException):
query = str(PostgreSQLQuery.into(self.table_abc).insert(1).on_conflict(self.table_abc.id))
def test_insert_on_conflict_two_handlers_do_update(self):
with self.assertRaises(QueryException):
query = PostgreSQLQuery.into(self.table_abc).insert(1).on_conflict(
self.table_abc.id).do_update(self.table_abc.name, "m").do_nothing()
def test_insert_returning_all_fields(self):
query = PostgreSQLQuery.into(self.table_abc).insert(1).returning(self.table_abc.star)
self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING *', str(query))
def test_insert_returning_from_other_table(self):
table_cba = Table('cba')
with self.assertRaises(QueryException):
PostgreSQLQuery.into(self.table_abc).insert(1).returning(table_cba.id)
def test_insert_returning_aggregate(self):
with self.assertRaises(QueryException):
PostgreSQLQuery.into(self.table_abc).insert(1).returning(Avg(self.table_abc.views))
def test_non_insert_on_conflict_do_update(self):
with self.assertRaises(QueryException):
query = PostgreSQLQuery.update(self.table_abc).set('foo', 'bar').on_conflict(
'id').do_update(['name'], ["m"])
def test_insert_all_columns_single_element_arrays(self):
query = PostgreSQLQuery.into(self.table_abc).insert((1, 'a', True)).returning(self.table_abc.star)
self.assertEqual('INSERT INTO "abc" VALUES (1,\'a\',true) RETURNING *', str(query))
def test_delete_returning_str(self):
q1 = PostgreSQLQuery.from_(self.table_abc).where(
self.table_abc.foo == self.table_abc.bar
).delete().returning('id')
self.assertEqual('DELETE FROM "abc" WHERE "foo"="bar" RETURNING id', str(q1))
def test_postgres_query_uses_double_quote_chars(self):
q = PostgreSQLQuery.from_(self.t).groupby(self.t.foo).select(self.t.foo)
self.assertEqual('SELECT "foo" FROM "abc" GROUP BY "foo"', str(q))
def test_insert_returning_null(self):
query = PostgreSQLQuery.into(self.table_abc).insert(1).returning(None)
self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING NULL', str(query))
def test_insert_returning_one_field_str(self):
query = PostgreSQLQuery.into(self.table_abc).insert(1).returning('id')
self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING id', str(query))
def test_insert_on_conflict_do_nothing_field_str(self):
query = PostgreSQLQuery.into(self.table_abc).insert(1).on_conflict('id').do_nothing()
self.assertEqual('INSERT INTO "abc" VALUES (1) ON CONFLICT (id) DO NOTHING', str(query))
def test_insert_returning_tuple(self):
query = PostgreSQLQuery.into(self.table_abc).insert(1).returning((1, 2, 3))
self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING (1,2,3)', str(query))
