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))