async def visit_order_by(self, query):
key = query._expr.value
table, field = key.split(".") if "." in key else (self.sql, key)
if type(table) == str:
table = Tables(table)[0]
self.sql_stack.append(lambda x: x.orderby(table.__getattr__(field)))
await self.visit(query.child)
async def visit_where(self, query):
# TODO: more general lazy expression evaluator
left, op, right = query._expr.value.split()
right = ast.literal_eval(right)
table, field = left.split(".") if "." in left else (self.sql, left)
if type(table) == str:
table = Tables(table)[0]
binary_op = _cmp_ops_dict[op]
self.sql_stack.append(
lambda x: x.where(binary_op(table.__getattr__(field), right)))
await self.visit(query.child)
def get_company_week_from_period(self, year: int, week: int,
company_id: str) -> int:
"""
Select company week of teh company
:param year: year
:param week: week
:param company_id: company id
:return: number of the week of the company in that given period
"""
survey_iterations, surveys, survey_questions, questions = Tables(
"`%s`.`%s`" % (self.__database, cons.SURVEYS_ITERATIONS_TABLE),
"`%s`.`%s`" % (self.__database, cons.SURVEYS_TABLE),
"`%s`.`%s`" % (self.__database, cons.SURVEYS_QUESTIONS_TABLE),
"`%s`.`%s`" % (self.__database, cons.QUESTIONS_TABLE))
q = Query.from_(survey_iterations) \
.join(surveys).on(surveys.id == survey_iterations.survey_id) \
.join(survey_questions).on(survey_questions.survey_iteration_id == survey_iterations.id) \
.join(questions).on(questions.id == survey_questions.question_id) \
.select(questions.week).distinct() \
.where((survey_iterations.week == week) & (survey_iterations.year == year)) \
.where(surveys.company_id == company_id) \
.where(questions.dimension_id != 1)
query = str(q).replace("\"", "")
result = self.__select_source(query)
return result.at[0, 0]
class DropTableTests(unittest.TestCase):
database_xyz = Database("mydb")
new_table, existing_table = Tables("abc", "efg")
foo, bar = Columns(("a", "INT"), ("b", "VARCHAR(100)"))
def test_drop_database(self):
q1 = Query.drop_database(self.database_xyz)
q2 = Query.drop_database(self.database_xyz).if_exists()
self.assertEqual('DROP DATABASE "mydb"', str(q1))
self.assertEqual('DROP DATABASE IF EXISTS "mydb"', str(q2))
def test_drop_table(self):
q1 = Query.drop_table(self.new_table)
q2 = Query.drop_table(self.new_table).if_exists()
self.assertEqual('DROP TABLE "abc"', str(q1))
self.assertEqual('DROP TABLE IF EXISTS "abc"', str(q2))
def test_drop_user(self):
q1 = Query.drop_user("myuser")
q2 = Query.drop_user("myuser").if_exists()
self.assertEqual('DROP USER "myuser"', str(q1))
self.assertEqual('DROP USER IF EXISTS "myuser"', str(q2))
def test_drop_view(self):
q1 = Query.drop_view("myview")
q2 = Query.drop_view("myview").if_exists()
self.assertEqual('DROP VIEW "myview"', str(q1))
self.assertEqual('DROP VIEW IF EXISTS "myview"', str(q2))
def new_game(self,
chat_id: int,
player1_id: int,
player2_id: int,
player1_name: str,
player2_name: str,
size_x: int,
size_y: int) -> None:
if not self.use_db:
return
games, players = Tables('games', 'players')
logging.info(f"Creating new game {chat_id} in DB")
player1_query = PostgreSQLQuery.into(players) \
.insert(player1_id, player1_name) \
.on_conflict(players.id).do_update(players.name, player1_name)
player2_query = PostgreSQLQuery.into(players) \
.insert(player2_id, player2_name) \
.on_conflict(players.id).do_update(players.name, player2_name)
new_game_query = PostgreSQLQuery.into(games) \
.insert(chat_id, player1_id, player2_id,
size_x, size_y, '{}', 'black',
player2_id, '', '', 'False,False')
self._cur.execute(player1_query.get_sql())
self._cur.execute(player2_query.get_sql())
self._cur.execute(new_game_query.get_sql())
self._conn.commit()
class SelectIntoTests(unittest.TestCase):
table_abc, table_efg, table_hij = Tables("abc", "efg", "hij")
def test_select_star_into(self):
query = Query.from_(self.table_abc).select("*").into(self.table_efg)
self.assertEqual('SELECT * INTO "efg" FROM "abc"', str(query))
def test_select_columns_into(self):
query = (Query.from_(
self.table_abc).select(self.table_abc.foo, self.table_abc.bar,
self.table_abc.buz).into(self.table_efg))
self.assertEqual('SELECT "foo","bar","buz" INTO "efg" FROM "abc"',
str(query))
def test_select_columns_into_with_join(self):
query = (Query.from_(self.table_abc).select(
self.table_abc.foo, self.table_abc.bar).join(self.table_hij).on(
self.table_abc.id == self.table_hij.abc_id).select(
self.table_hij.fiz,
self.table_hij.buz).into(self.table_efg))
self.assertEqual(
'SELECT "abc"."foo","abc"."bar","hij"."fiz","hij"."buz" '
'INTO "efg" FROM "abc" '
'JOIN "hij" ON "abc"."id"="hij"."abc_id"',
str(query),
)
def test_union_with_no_quote_char(self):
abc, efg = Tables('abc', 'efg')
hij = Query.from_(abc).select(abc.t).union(Query.from_(efg).select(efg.t))
q = Query.from_(hij).select(fn.Avg(hij.t))
self.assertEqual('SELECT AVG(sq0.t) FROM ((SELECT t FROM abc) UNION (SELECT t FROM efg)) sq0',
q.get_sql(quote_char=None))
def test_cross_join_on_subquery(self):
table_a, table_b = Tables("a", "b")
q_a = Query.from_(table_a).select("*")
q_b = Query.from_(table_b).select("*").join(q_a).cross().select("*")
self.assertEqual(
'SELECT * FROM "b" CROSS JOIN (SELECT * FROM "a") "sq0"', str(q_b))
def get_stocklist(date, engine, dategap=90):
'''
Get stocklist that :
1) listed before dategap from date, and has not been delisted till date
b.Delisting == '19000101' means still on trade
2) on date, Turnover > 0
Params:
date:
string, like '%Y%m%d'
engine:
database engine
dategap:
days from startdate to date
Returns:
list, list of stockcode like ['000001.SZ', '000002.SZ', ...]
'''
startdate = (pd.to_datetime(date) -
timedelta(days=dategap)).strftime('%Y%m%d')
a, b = Tables('stock_alpha_factors', 'stock_listing_info')
q = Query.from_(a).join(b).on(a.stock_id == b.stock_id).select(
a.stock_id).where((a.trade_date == date) & (a.Turnover > 0)
& (b.Listing <= startdate)
& ((b.Delisting > date)
| (b.Delisting == '19000101'))).orderby(
a.stock_id)
stocklist = (pd.read_sql(q.get_sql(),
engine).stock_id.drop_duplicates().tolist())
return stocklist
class UnionTests(unittest.TestCase):
table1, table2 = Tables('abc', 'efg')
def test_union(self):
query1 = Query.from_(self.table1).select(self.table1.foo)
query2 = Query.from_(self.table2).select(self.table2.bar)
self.assertEqual(
'SELECT "foo" FROM "abc" UNION SELECT "bar" FROM "efg"',
str(query1 + query2))
self.assertEqual(
'SELECT "foo" FROM "abc" UNION SELECT "bar" FROM "efg"',
str(query1.union(query2)))
def test_union_all(self):
query1 = Query.from_(self.table1).select(self.table1.foo)
query2 = Query.from_(self.table2).select(self.table2.bar)
self.assertEqual(
'SELECT "foo" FROM "abc" UNION ALL SELECT "bar" FROM "efg"',
str(query1 * query2))
self.assertEqual(
'SELECT "foo" FROM "abc" UNION ALL SELECT "bar" FROM "efg"',
str(query1.union_all(query2)))
def test_require_equal_number_of_fields(self):
query1 = Query.from_(self.table1).select(self.table1.foo)
query2 = Query.from_(self.table2).select(self.table2.fiz,
self.table2.buz)
with self.assertRaises(UnionException):
str(query1 + query2)
class ArrayTests(unittest.TestCase):
table_abc, table_efg = Tables("abc", "efg")
def test_array_general(self):
query = Query.from_(self.table_abc).select(Array(1, "a", ["b", 2, 3]))
self.assertEqual("SELECT [1,'a',['b',2,3]] FROM \"abc\"", str(query))
def test_empty_psql_array(self):
query = PostgreSQLQuery.from_(self.table_abc).select(Array())
self.assertEqual("SELECT '{}' FROM \"abc\"", str(query))
def test_psql_array_general(self):
query = PostgreSQLQuery.from_(self.table_abc).select(
Array(1, Array(2, 2, 2), 3))
self.assertEqual('SELECT ARRAY[1,ARRAY[2,2,2],3] FROM "abc"',
str(query))
def test_render_alias_in_array_sql(self):
tb = Table("tb")
q = Query.from_(tb).select(Array(tb.col).as_("different_name"))
self.assertEqual(str(q), 'SELECT ["col"] "different_name" FROM "tb"')
class ImmutabilityTests(unittest.TestCase):
table_a, table_b = Tables("a", "b")
def test_select_returns_new_query_instance(self):
query_a = Query.from_(self.table_a).select(self.table_a.foo)
query_b = query_a.select(self.table_a.bar)
self.assertIsNot(query_a, query_b)
self.assertNotEqual(str(query_a), str(query_b))
def test_queries_after_join(self):
query1 = Query.from_(self.table_a).select(self.table_a.foo)
query2 = (query1.join(self.table_b).on(
self.table_a.foo == self.table_b.bar).select(self.table_b.buz))
self.assertEqual('SELECT "foo" FROM "a"', str(query1))
self.assertEqual(
'SELECT "a"."foo","b"."buz" FROM "a" JOIN "b" ON "a"."foo"="b"."bar"',
str(query2),
)
def test_immutable_kwarg_on_query_builder_disables_immutability(self):
query0 = Query.from_(self.table_a, immutable=False)
query1 = query0.select(self.table_a.foo)
self.assertIs(query0, query1)
def test_union_as_subquery(self):
abc, efg = Tables('abc', 'efg')
hij = Query.from_(abc).select(abc.t).union(Query.from_(efg).select(efg.t))
q = Query.from_(hij).select(fn.Avg(hij.t))
self.assertEqual('SELECT AVG("sq0"."t") FROM ((SELECT "t" FROM "abc") UNION (SELECT "t" FROM "efg")) "sq0"',
str(q))
def test_join_table_on_update_query(self):
a, b = Tables("a", "b")
q = Query.update(a).join(b).on(a.fkey_id == b.id).where(b.foo == 1).set("adwords_batch_job_id", 1)
self.assertEqual(
'UPDATE "a" ' 'JOIN "b" ' 'ON "a"."fkey_id"="b"."id" ' 'SET "adwords_batch_job_id"=1 ' 'WHERE "b"."foo"=1',
str(q),
)
def test_project(self):
sql = (UserQuery(range(1, 10)).project([":id", "name"]).where(
ast.Expr('user.age >= 16')).order_by(
ast.Expr('user.age')).take(3).to_sql())
user = Tables("user")[0]
expected = (Query.from_("user").select(
"id", "name").where(user.age >= 16).orderby(user.age).limit(3))
self.assertEqual(str(expected), str(sql))
def test_join_second_table_in_from_clause(self):
table_a, table_b, table_c = Tables("a", "b", "c")
q = Query.from_(table_a).from_(table_b).select('*') \
.join(table_c).on(table_b.c_id == table_c.id)
self.assertEqual('SELECT * '
'FROM "a","b" '
'JOIN "c" ON "b"."c_id"="c"."id"', str(q))
def test_many_tables_with_or_not_alias(self):
tables_data = [('table1', 't1'), ('table2'), 'table3']
tables = Tables(*tables_data)
for i in range(len(tables)):
if isinstance(tables_data[i], tuple):
self.assertIsNotNone(tables[i].alias)
else:
self.assertIsNone(tables[i].alias)
def test_many_tables_with_or_not_alias(self):
tables_data = [("table1", "t1"), ("table2"), "table3"]
tables = Tables(*tables_data)
for i in range(len(tables)):
if isinstance(tables_data[i], tuple):
self.assertIsNotNone(tables[i].alias)
else:
self.assertIsNone(tables[i].alias)
def search(key):
room, usr = Tables(cfg.room, cfg.usr)
q = Query.from_(room).select(
room.id, room.name,
fn.Concat(1).as_('type')).where(
room.name.like(key + '%')) + Query.from_(usr).select(
usr.id, usr.name,
fn.Concat(0).as_('type')).where(usr.name.like(key + '%'))
return q
def test_intersect_with_no_quote_char(self):
abc, efg = Tables("abc", "efg")
hij = Query.from_(abc).select(abc.t).intersect(Query.from_(efg).select(efg.t))
q = Query.from_(hij).select(fn.Avg(hij.t))
self.assertEqual(
"SELECT AVG(sq0.t) FROM ((SELECT t FROM abc) INTERSECT (SELECT t FROM efg)) sq0",
q.get_sql(quote_char=None),
)
def test_intersect_as_subquery(self):
abc, efg = Tables("abc", "efg")
hij = Query.from_(abc).select(abc.t).intersect(Query.from_(efg).select(efg.t))
q = Query.from_(hij).select(fn.Avg(hij.t))
self.assertEqual(
'SELECT AVG("sq0"."t") FROM ((SELECT "t" FROM "abc") INTERSECT (SELECT "t" FROM "efg")) "sq0"',
str(q),
)
class QueryTests(unittest.TestCase):
manager = QueryManager()
maxDiff = None
mock_table, mock_join1, mock_join2 = Tables('test_table', 'test_join1', 'test_join2')
@classmethod
def setUpClass(cls):
settings.database = TestDatabase()
def test_join_on_collate(self):
table_a, table_b = Tables("a", "b")
q1 = Query.from_(table_a).select(table_b.ouch).join(table_b).on(table_a.foo == table_b.boo,
collate="utf8_general_ci")
q2 = Query.from_(table_a).select(table_b.ouch).join(table_b).on(table_a.foo == table_b.boo)
self.assertEqual('SELECT "b"."ouch" FROM "a" JOIN "b" ON "a"."foo"="b"."boo" COLLATE utf8_general_ci', str(q1))
self.assertEqual('SELECT "b"."ouch" FROM "a" JOIN "b" ON "a"."foo"="b"."boo"', str(q2))
class BracketTests(unittest.TestCase):
table_abc, table_efg = Tables('abc', 'efg')
def test_arithmetic_with_brackets(self):
q = Query \
.from_(self.table_abc) \
.select(Bracket(self.table_abc.foo / 2) / 2)
self.assertEqual('SELECT ("foo"/2)/2 FROM "abc"', str(q))
class InsertSelectFromTests(unittest.TestCase):
table_abc, table_efg, table_hij = Tables('abc', 'efg', 'hij')
def test_insert_star(self):
query = Query.into(self.table_abc).from_(self.table_efg).select('*')
self.assertEqual('INSERT INTO "abc" SELECT * FROM "efg"', str(query))
def test_insert_ignore_star(self):
query = Query.into(self.table_abc).from_(
self.table_efg).select('*').ignore()
self.assertEqual('INSERT IGNORE INTO "abc" SELECT * FROM "efg"',
str(query))
def test_insert_from_columns(self):
query = Query.into(self.table_abc).from_(self.table_efg).select(
self.table_efg.fiz, self.table_efg.buz, self.table_efg.baz)
self.assertEqual(
'INSERT INTO "abc" '
'SELECT "fiz","buz","baz" FROM "efg"', str(query))
def test_insert_columns_from_star(self):
query = Query.into(self.table_abc).columns(
self.table_abc.foo,
self.table_abc.bar,
self.table_abc.buz,
).from_(self.table_efg).select('*')
self.assertEqual(
'INSERT INTO "abc" ("foo","bar","buz") '
'SELECT * FROM "efg"', str(query))
def test_insert_columns_from_columns(self):
query = Query.into(self.table_abc).columns(
self.table_abc.foo, self.table_abc.bar,
self.table_abc.buz).from_(self.table_efg).select(
self.table_efg.fiz, self.table_efg.buz, self.table_efg.baz)
self.assertEqual(
'INSERT INTO "abc" ("foo","bar","buz") '
'SELECT "fiz","buz","baz" FROM "efg"', str(query))
def test_insert_columns_from_columns_with_join(self):
query = Query.into(self.table_abc).columns(
self.table_abc.c1, self.table_abc.c2, self.table_abc.c3,
self.table_abc.c4).from_(self.table_efg).select(
self.table_efg.foo,
self.table_efg.bar).join(self.table_hij).on(
self.table_efg.id == self.table_hij.abc_id).select(
self.table_hij.fiz, self.table_hij.buz)
self.assertEqual(
'INSERT INTO "abc" ("c1","c2","c3","c4") '
'SELECT "efg"."foo","efg"."bar","hij"."fiz","hij"."buz" FROM "efg" '
'JOIN "hij" ON "efg"."id"="hij"."abc_id"', str(query))
class QuoteTests(unittest.TestCase):
maxDiff = None
table_abc, table_efg = Tables("abc", "efg")
def setUp(self):
subquery1 = (Query.from_(self.table_abc).select(
self.table_abc.foo,
fn.Sum(self.table_abc.fizz + self.table_abc.buzz).as_("fizzbuzz"),
).groupby(self.table_abc.foo))
subquery2 = Query.from_(self.table_efg).select(
self.table_efg.foo.as_("foo_two"),
self.table_efg.bar,
)
self.query = (Query.from_(subquery1).select(
subquery1.foo, subquery1.fizzbuzz).join(subquery2).on(
subquery1.foo == subquery2.foo_two).select(
subquery2.foo_two, subquery2.bar))
def test_replace_quote_char_in_complex_query(self):
self.assertEqual(
"SELECT "
"`sq0`.`foo`,`sq0`.`fizzbuzz`,"
"`sq1`.`foo_two`,`sq1`.`bar` "
"FROM ("
"SELECT "
"`foo`,SUM(`fizz`+`buzz`) `fizzbuzz` "
"FROM `abc` "
"GROUP BY `foo`"
") `sq0` JOIN ("
"SELECT "
"`foo` `foo_two`,`bar` "
"FROM `efg`"
") `sq1` ON `sq0`.`foo`=`sq1`.`foo_two`",
self.query.get_sql(quote_char="`"),
)
def test_no_quote_char_in_complex_query(self):
self.assertEqual(
"SELECT "
"sq0.foo,sq0.fizzbuzz,"
"sq1.foo_two,sq1.bar "
"FROM ("
"SELECT "
"foo,SUM(fizz+buzz) fizzbuzz "
"FROM abc "
"GROUP BY foo"
") sq0 JOIN ("
"SELECT "
"foo foo_two,bar "
"FROM efg"
") sq1 ON sq0.foo=sq1.foo_two",
self.query.get_sql(quote_char=None),
)
def test_join_on_subquery_criteria(self):
table_a, table_b, table_c = Tables('a', 'b', 'c')
subquery = Query.from_(table_c).select('id').limit(1)
query = Query.from_(table_a).select('*') \
.join(table_b).on((table_a.b_id == table_b.id) & (table_b.c_id == subquery))
self.assertEqual('SELECT * '
'FROM "a" '
'JOIN "b" ON "a"."b_id"="b"."id" AND "b"."c_id"='
'(SELECT "id" FROM "c" LIMIT 1)', str(query))
def get_column_definitions(self, schema, table, connection=None):
view_columns, table_columns = Tables('view_columns', 'columns')
view_query = (VerticaQuery.from_(view_columns).select(
view_columns.column_name, view_columns.data_type).where(
(view_columns.table_schema == Parameter(':schema'))
& (view_columns.field('table_name') == Parameter(':table'))).
distinct())
table_query = (VerticaQuery.from_(
table_columns, immutable=False).select(
table_columns.column_name, table_columns.data_type).where(
(table_columns.table_schema == Parameter(':schema'))
& (table_columns.field("table_name") == Parameter(':table')
)).distinct())
return self.fetch(str(view_query + table_query),
parameters=dict(schema=schema, table=table),
connection=connection)
class CreateTemporaryTableTests(unittest.TestCase):
new_table, existing_table = Tables('abc', 'efg')
foo, bar = Columns(('a', 'INT'), ('b', 'VARCHAR(100)'))
select = VerticaQuery.from_(existing_table).select(existing_table.foo, existing_table.bar)
def test_create_local_temporary_table(self):
with self.subTest('with columns'):
q = VerticaQuery \
.create_table(self.new_table) \
.temporary() \
.local() \
.columns(self.foo, self.bar)
self.assertEqual('CREATE LOCAL TEMPORARY TABLE "abc" ("a" INT,"b" VARCHAR(100))', str(q))
with self.subTest('with select'):
q = VerticaQuery \
.create_table(self.new_table) \
.temporary() \
.local() \
.as_select(self.select)
self.assertEqual('CREATE LOCAL TEMPORARY TABLE "abc" AS (SELECT "foo","bar" FROM "efg")', str(q))
def test_create_local_table_without_temporary_raises_error(self):
with self.assertRaises(AttributeError):
VerticaQuery \
.create_table(self.new_table) \
.local()
def test_create_temporary_table_preserve_rows(self):
with self.subTest('with columns'):
q = VerticaQuery \
.create_table(self.new_table) \
.temporary() \
.preserve_rows() \
.columns(self.foo, self.bar)
self.assertEqual('CREATE TEMPORARY TABLE "abc" ("a" INT,"b" VARCHAR(100)) ON COMMIT PRESERVE ROWS', str(q))
with self.subTest('with select'):
q = VerticaQuery \
.create_table(self.new_table) \
.temporary() \
.preserve_rows() \
.as_select(self.select)
self.assertEqual('CREATE TEMPORARY TABLE "abc" ON COMMIT PRESERVE ROWS AS (SELECT "foo","bar" FROM "efg")',
str(q))
def test_create_table_preserve_rows_without_temporary_raises_error(self):
with self.assertRaises(AttributeError):
VerticaQuery \
.create_table(self.new_table) \
.preserve_rows()
class HavingTests(unittest.TestCase):
table_abc, table_efg = Tables("abc", "efg")
def test_having_greater_than(self):
q = (Query.from_(self.table_abc).select(
self.table_abc.foo, fn.Sum(self.table_abc.bar)).groupby(
self.table_abc.foo).having(fn.Sum(self.table_abc.bar) > 1))
self.assertEqual(
'SELECT "foo",SUM("bar") FROM "abc" GROUP BY "foo" HAVING SUM("bar")>1',
str(q),
)
def test_having_and(self):
q = (Query.from_(self.table_abc).select(
self.table_abc.foo, fn.Sum(self.table_abc.bar)).groupby(
self.table_abc.foo).having((fn.Sum(self.table_abc.bar) > 1) &
(fn.Sum(self.table_abc.bar) < 100)))
self.assertEqual(
'SELECT "foo",SUM("bar") FROM "abc" GROUP BY "foo" HAVING SUM("bar")>1 AND SUM("bar")<100',
str(q),
)
def test_having_join_and_equality(self):
q = (Query.from_(self.table_abc).join(self.table_efg).on(
self.table_abc.foo == self.table_efg.foo).select(
self.table_abc.foo, fn.Sum(self.table_efg.bar),
self.table_abc.buz).groupby(self.table_abc.foo).having(
self.table_abc.buz == "fiz").having(
fn.Sum(self.table_efg.bar) > 100))
self.assertEqual(
'SELECT "abc"."foo",SUM("efg"."bar"),"abc"."buz" FROM "abc" '
'JOIN "efg" ON "abc"."foo"="efg"."foo" '
'GROUP BY "abc"."foo" '
'HAVING "abc"."buz"=\'fiz\' AND SUM("efg"."bar")>100',
str(q),
)
def test_mysql_query_uses_backtick_quote_chars(self):
q = (MySQLQuery.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_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))
