async def test_schema(self):
from asyncpg.exceptions import InvalidSchemaNameError
self.db_config["connections"]["models"]["credentials"][
"schema"] = "mytestschema"
await Tortoise.init(self.db_config, _create_db=True)
with self.assertRaises(InvalidSchemaNameError):
await Tortoise.generate_schemas()
conn = Tortoise.get_connection("models")
await conn.execute_script("CREATE SCHEMA mytestschema;")
await Tortoise.generate_schemas()
tournament = await Tournament.create(name="Test")
await Tortoise.close_connections()
del self.db_config["connections"]["models"]["credentials"]["schema"]
await Tortoise.init(self.db_config)
with self.assertRaises(OperationalError):
await Tournament.filter(name="Test").first()
conn = Tortoise.get_connection("models")
res = await conn.execute_query(
"SELECT id, name FROM mytestschema.tournament WHERE name='Test' LIMIT 1"
)
self.assertEqual(len(res), 1)
self.assertEqual(tournament.id, res[0][0])
self.assertEqual(tournament.name, res[0][1])
async def run():
await Tortoise.init(
{
"connections": {
"first": 'mysql://*****:*****@localhost:55555/test_demo',
"second": 'mysql://*****:*****@localhost:55555/test_demo2',
},
"apps": {
"tournaments": {"models": ["__main__"], "default_connection": "first"},
"events": {"models": ["__main__"], "default_connection": "second"},
},
}
)
# await Tortoise.generate_schemas()
client = Tortoise.get_connection("first")
second_client = Tortoise.get_connection("second")
tournament = await Tournament.create(name="Tournament")
await Event(name="Event", tournament_id=tournament.id).save()
try:
# await client.execute_query('SELECT * FROM event')
print(await Tournament.all())
except OperationalError:
print("Expected it to fail")
results = await second_client.execute_query('SELECT * FROM event')
print(results)
event = await Event.filter(id=1).first()
if event:
team = Team(name='cjs')
await team.save()
await event.participants.add(team)
print(await event.participants.all())
async def setUp(self):
if Tortoise._inited:
await self._tearDownDB()
first_db_config = test.getDBConfig(
app_label='models',
modules=['tortoise.tests.testmodels'],
)
second_db_config = test.getDBConfig(
app_label='events',
modules=['tortoise.tests.testmodels'],
)
merged_config = {
'connections': {
**first_db_config['connections'],
**second_db_config['connections']
},
'apps': {
**first_db_config['apps'],
**second_db_config['apps']
},
}
await Tortoise.init(merged_config, _create_db=True)
await Tortoise.generate_schemas()
self.db = Tortoise.get_connection('models')
self.second_db = Tortoise.get_connection('events')
async def run():
await Tortoise.init(
{
"connections": {
"first": {
"engine": "tortoise.backends.sqlite",
"credentials": {"file_path": "example.sqlite3"},
},
"second": {
"engine": "tortoise.backends.sqlite",
"credentials": {"file_path": "example1.sqlite3"},
},
},
"apps": {
"tournaments": {"models": ["__main__"], "default_connection": "first"},
"events": {"models": ["__main__"], "default_connection": "second"},
},
}
)
await Tortoise.generate_schemas()
client = Tortoise.get_connection("first")
second_client = Tortoise.get_connection("second")
tournament = await Tournament.create(name="Tournament")
await Event(name="Event", tournament_id=tournament.id).save()
try:
await client.execute_query('SELECT * FROM "event"')
except OperationalError:
print("Expected it to fail")
results = await second_client.execute_query('SELECT * FROM "event"')
print(results)
async def get_db_status(self):
read = Tortoise.get_connection('read')
write = Tortoise.get_connection('write')
return {
'read_connection': await self.check_connection(read),
'write_connection': await self.check_connection(write)
}
async def asyncSetUp(self):
await super().asyncSetUp()
if Tortoise._inited:
await self._tearDownDB()
first_db_config = test.getDBConfig(app_label="models", modules=["tests.testmodels"])
second_db_config = test.getDBConfig(app_label="events", modules=["tests.testmodels"])
merged_config = {
"connections": {**first_db_config["connections"], **second_db_config["connections"]},
"apps": {**first_db_config["apps"], **second_db_config["apps"]},
}
await Tortoise.init(merged_config, _create_db=True)
await Tortoise.generate_schemas()
self.db = Tortoise.get_connection("models")
self.second_db = Tortoise.get_connection("events")
async def initialize_tests(event_loop, request):
await Tortoise.init(config=tortoise_orm, _create_db=True)
await generate_schema_for_client(Tortoise.get_connection("default"),
safe=True)
client = Tortoise.get_connection("default")
if client.schema_generator is MySQLSchemaGenerator:
Migrate.ddl = MysqlDDL(client)
elif client.schema_generator is SqliteSchemaGenerator:
Migrate.ddl = SqliteDDL(client)
elif client.schema_generator is AsyncpgSchemaGenerator:
Migrate.ddl = PostgresDDL(client)
Migrate.dialect = Migrate.ddl.DIALECT
request.addfinalizer(
lambda: event_loop.run_until_complete(Tortoise._drop_databases()))
async def test():
await Tortoise.init(db_url=DB_URL,
modules={'models': ['Utils.DataModels']})
# Generate the schema
await Tortoise.generate_schemas()
connection = Tortoise.get_connection("default")
await connection.execute_script("alter table apitokens modify id BIGINT")
async def _migrate_for_v30() -> None:
await init_db(create_db=False)
client = Tortoise.get_connection('default')
await client.execute_script(
'''CREATE TABLE `activity` (
`can_comment` tinyint(1) NOT NULL,
`id` int(11) NOT NULL AUTO_INCREMENT,
`created_at` datetime(6) NOT NULL,
`user_id` int(11) NOT NULL,
`target_id` int(11) NOT NULL,
`target_kind` int(11) NOT NULL,
PRIMARY KEY (`id`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4''')
await client.execute_script(
'''CREATE TABLE `statuses` (
`id` int(11) NOT NULL AUTO_INCREMENT,
`created_at` datetime(6) NOT NULL,
`user_id` int(11) NOT NULL,
PRIMARY KEY (`id`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4''')
await client.execute_script(
'ALTER TABLE comments CHANGE `post_id` `target_id` int(11) NOT NULL')
await client.execute_script(
'alter table comments add column `target_kind` smallint(6) DEFAULT 1001')
async def run():
await Tortoise.init(db_url="sqlite://:memory:",
modules={"models": ["__main__"]})
await Tortoise.generate_schemas()
# Need to get a connection. Unless explicitly specified, the name should be 'default'
conn = Tortoise.get_connection("default")
# Now we can execute queries in the normal autocommit mode
await conn.execute_query("INSERT INTO event (name) VALUES ('Foo')")
# You can also you parameters, but you need to use the right param strings for each dialect
await conn.execute_query("INSERT INTO event (name) VALUES (?)", ["Bar"])
# To do a transaction you'd need to use the in_transaction context manager
async with in_transaction("default") as tconn:
await tconn.execute_query("INSERT INTO event (name) VALUES ('Moo')")
# Unless an exception happens it should commit automatically
# This transaction is rolled back
async with in_transaction("default") as tconn:
await tconn.execute_query("INSERT INTO event (name) VALUES ('Sheep')")
# Rollback to fail transaction
await tconn.rollback()
# Consider using execute_query_dict to get return values as a dict
val = await conn.execute_query_dict("SELECT * FROM event")
print(val)
async def init() -> None:
await init_db(create_db=False)
await Tortoise._drop_databases()
await init_db(create_db=True)
await Tortoise.generate_schemas()
# Add Indexes
client = Tortoise.get_connection('default')
await client.execute_script(
'alter table posts add index `idx_slug` (`slug`)')
await client.execute_script(
'alter table post_tags add index `idx_post_tag` (`post_id`, `tag_id`)')
await client.execute_script(
'alter table comments add index `idx_target_kind` (`target_id`, `target_kind`)'
)
await client.execute_script(
'alter table react_items add index `idx_id_kind_user` (`target_id`, `target_kind`, `user_id`)'
) # noqa
if not await client.execute_query(
'show columns from `posts` like "pageview"'):
await migrate_for_v25()
try:
await migrate_for_v35()
except OperationalError:
...
async def run():
await Tortoise.init(db_url="mysql://*****:*****@localhost:55555/test_demo", modules={"models": ["__main__"]})
client = Tortoise.get_connection('default')
result = await client.execute_query("SELECT * FROM team")
print(result)
async def resolve_places_by_range(cls, info: ResolveInfo, latitude: float,
longitude: float,
distance: float) -> List[PlaceModel]:
distance *= 1.05
connection = Tortoise.get_connection('models')
raw_data: List[dict] = await connection.execute_query_dict(
"""
SELECT * FROM (
SELECT id,
created_at,
updated_at,
address,
longitude,
latitude,
work_time_start,
work_time_stop,
preorder,
rating,
min_intervals_for_book,
max_intervals_for_book,
restaurant_id,
(6371 * ACOS(SIN(RADIANS($1)) * SIN(RADIANS(latitude)) +
COS(RADIANS($1)) * COS(RADIANS(latitude)) *
(COS(RADIANS($2) - RADIANS(longitude))))) / 2 as distance
FROM place
ORDER BY distance
) subquery
WHERE distance < $3;
""", [latitude, longitude, distance])
places = [PlaceModel(**data) for data in raw_data]
return places
async def asyncSetUp(self):
await super().asyncSetUp()
# Build large dataset
self.intfields = [
await IntFields.create(intnum=val) for val in range(10, 100, 3)
]
self.db = Tortoise.get_connection("models")
async def get_strategy_code(task_id: str, request: Request):
"""检查是否需要因此策略信息"""
if not await check_task_permission(TaskType.PAPER_TRADING, task_id, request):
raise HTTPException(status.HTTP_403_FORBIDDEN, "没有权限")
# get code
query_str = "SELECT backtest_id FROM wk_simulation WHERE task_id=%s"
client = Tortoise.get_connection("qpweb")
try:
rows = await client.execute_query_dict(query_str, task_id)
except TypeError:
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="复制错误，策略已不存在")
if len(rows) != 1:
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="复制错误，策略已不存在")
backtest_id = rows[0]["backtest_id"]
query_str = f"SELECT code FROM wk_strategy_backtest WHERE id=%s"
try:
rows = await client.execute_query_dict(query_str, backtest_id)
except TypeError:
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="复制错误，策略已不存在")
if len(rows) != 1:
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="复制错误，策略已不存在")
code = rows[0]["code"]
return CommonOut(data=code)
def skip_wrapper(*args, **kwargs):
db = Tortoise.get_connection(connection_name)
for key, val in conditions.items():
if getattr(db.capabilities, key) != val:
raise SkipTest("Capability {key} != {val}".format(key=key,
val=val))
return test_item(*args, **kwargs)
async def test_shorthand_init(self):
await Tortoise.init(
db_url="sqlite://{}".format(":memory:"),
modules={"models": ["tortoise.tests.testmodels"]},
)
self.assertIn("models", Tortoise.apps)
self.assertIsNotNone(Tortoise.get_connection("default"))
async def leaderboard_raw(network_id,
from_date=0,
to_date=0,
offset=0,
limit=10,
currency='rune'):
if currency == 'rune':
sum_variable = "rune_volume"
else:
sum_variable = "usd_volume"
end_date_cond = f' AND date <= {int(to_date)} ' if to_date else ''
q = (f"SELECT "
f" user_address,"
f" SUM({sum_variable}) total_volume, MAX(date) as date, "
f" COUNT(id) n "
f" FROM thortx "
f" WHERE network = '{network_id}' "
f" AND type = '{ThorTxType.TYPE_SWAP}' "
f' AND date >= {int(from_date)} {end_date_cond}'
f" GROUP BY user_address "
f" ORDER BY SUM({sum_variable}) "
f" DESC LIMIT {int(limit)} OFFSET {int(offset)}")
conn = Tortoise.get_connection("default")
return await conn.execute_query_dict(q)
async def init_db(args):
await Tortoise.init(config=settings.TORTOISE_ORM)
if args.test:
sql = get_schema_sql(Tortoise.get_connection('default'), safe=False)
print(sql)
else:
await Tortoise.generate_schemas()
async def test_schema_safe(self):
self.maxDiff = None
await self.init_for("tests.models_schema_create")
sql = get_schema_sql(Tortoise.get_connection("default"), safe=True)
self.assertEqual(
sql.strip(),
"""
CREATE TABLE IF NOT EXISTS "defaultpk" (
"id" SERIAL NOT NULL PRIMARY KEY,
"val" INT NOT NULL
);
CREATE TABLE IF NOT EXISTS "sometable" (
"sometable_id" SERIAL NOT NULL PRIMARY KEY,
"some_chars_table" VARCHAR(255) NOT NULL,
"fk_sometable" INT REFERENCES "sometable" (sometable_id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS "sometable_some_ch_115115_idx" ON "sometable" (some_chars_table);
CREATE TABLE IF NOT EXISTS "team" (
"name" VARCHAR(50) NOT NULL PRIMARY KEY,
"manager_id" VARCHAR(50) REFERENCES "team" (name) ON DELETE CASCADE
);
COMMENT ON COLUMN team.name IS 'The TEAM name (and PK)';
COMMENT ON TABLE team IS 'The TEAMS!';
CREATE TABLE IF NOT EXISTS "tournament" (
"tid" SMALLSERIAL NOT NULL PRIMARY KEY,
"name" TEXT NOT NULL,
"created" TIMESTAMP NOT NULL
);
CREATE INDEX IF NOT EXISTS "tournament_name_116110_idx" ON "tournament" (name);
COMMENT ON COLUMN tournament.name IS 'Tournament name';
COMMENT ON COLUMN tournament.created IS 'Created */''`/* datetime';
COMMENT ON TABLE tournament IS 'What Tournaments */''`/* we have';
CREATE TABLE IF NOT EXISTS "event" (
"id" BIGSERIAL NOT NULL PRIMARY KEY,
"name" TEXT NOT NULL UNIQUE,
"modified" TIMESTAMP NOT NULL,
"prize" DECIMAL(10,2),
"token" VARCHAR(100) NOT NULL UNIQUE,
"tournament_id" SMALLINT NOT NULL REFERENCES "tournament" (tid) ON DELETE CASCADE
);
COMMENT ON COLUMN event.id IS 'Event ID';
COMMENT ON COLUMN event.token IS 'Unique token';
COMMENT ON COLUMN event.tournament_id IS 'FK to tournament';
COMMENT ON TABLE event IS 'This table contains a list of all the events';
CREATE TABLE IF NOT EXISTS "sometable_self" (
"backward_sts" INT NOT NULL REFERENCES "sometable" (sometable_id) ON DELETE CASCADE,
"sts_forward" INT NOT NULL REFERENCES "sometable" (sometable_id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS "team_team" (
"team_rel_id" VARCHAR(50) NOT NULL REFERENCES "team" (name) ON DELETE CASCADE,
"team_id" VARCHAR(50) NOT NULL REFERENCES "team" (name) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS "teamevents" (
"event_id" BIGINT NOT NULL REFERENCES "event" (id) ON DELETE CASCADE,
"team_id" VARCHAR(50) NOT NULL REFERENCES "team" (name) ON DELETE CASCADE
);
COMMENT ON TABLE teamevents IS 'How participants relate';
""".strip(),
)
async def test_schema_safe(self):
self.maxDiff = None
await self.init_for("tests.models_schema_create")
sql = get_schema_sql(Tortoise.get_connection("default"), safe=True)
self.assertEqual(
sql.strip(),
"""
CREATE TABLE IF NOT EXISTS `defaultpk` (
`id` INT NOT NULL PRIMARY KEY AUTO_INCREMENT,
`val` INT NOT NULL
) CHARACTER SET utf8mb4;
CREATE TABLE IF NOT EXISTS `sometable` (
`sometable_id` INT NOT NULL PRIMARY KEY AUTO_INCREMENT,
`some_chars_table` VARCHAR(255) NOT NULL,
`fk_sometable` INT,
CONSTRAINT `fk_sometabl_sometabl_6efae9bd` FOREIGN KEY (`fk_sometable`) REFERENCES `sometable` (`sometable_id`) ON DELETE CASCADE,
KEY `sometable_some_ch_115115_idx` (`some_chars_table`)
) CHARACTER SET utf8mb4;
CREATE TABLE IF NOT EXISTS `team` (
`name` VARCHAR(50) NOT NULL PRIMARY KEY COMMENT 'The TEAM name (and PK)',
`manager_id` VARCHAR(50),
CONSTRAINT `fk_team_team_9c77cd8f` FOREIGN KEY (`manager_id`) REFERENCES `team` (`name`) ON DELETE CASCADE
) CHARACTER SET utf8mb4 COMMENT='The TEAMS!';
CREATE TABLE IF NOT EXISTS `tournament` (
`tid` SMALLINT NOT NULL PRIMARY KEY AUTO_INCREMENT,
`name` VARCHAR(100) NOT NULL COMMENT 'Tournament name',
`created` DATETIME(6) NOT NULL COMMENT 'Created */\\'`/* datetime',
KEY `tournament_name_116110_idx` (`name`)
) CHARACTER SET utf8mb4 COMMENT='What Tournaments */\\'`/* we have';
CREATE TABLE IF NOT EXISTS `event` (
`id` BIGINT NOT NULL PRIMARY KEY AUTO_INCREMENT COMMENT 'Event ID',
`name` TEXT NOT NULL,
`modified` DATETIME(6) NOT NULL,
`prize` DECIMAL(10,2),
`token` VARCHAR(100) NOT NULL UNIQUE COMMENT 'Unique token',
`tournament_id` SMALLINT NOT NULL COMMENT 'FK to tournament',
CONSTRAINT `fk_event_tourname_51c2b82d` FOREIGN KEY (`tournament_id`) REFERENCES `tournament` (`tid`) ON DELETE CASCADE
) CHARACTER SET utf8mb4 COMMENT='This table contains a list of all the events';
CREATE TABLE IF NOT EXISTS `sometable_self` (
`backward_sts` INT NOT NULL,
`sts_forward` INT NOT NULL,
FOREIGN KEY (`backward_sts`) REFERENCES `sometable` (`sometable_id`) ON DELETE CASCADE,
FOREIGN KEY (`sts_forward`) REFERENCES `sometable` (`sometable_id`) ON DELETE CASCADE
) CHARACTER SET utf8mb4;
CREATE TABLE IF NOT EXISTS `team_team` (
`team_rel_id` VARCHAR(50) NOT NULL,
`team_id` VARCHAR(50) NOT NULL,
FOREIGN KEY (`team_rel_id`) REFERENCES `team` (`name`) ON DELETE CASCADE,
FOREIGN KEY (`team_id`) REFERENCES `team` (`name`) ON DELETE CASCADE
) CHARACTER SET utf8mb4;
CREATE TABLE IF NOT EXISTS `teamevents` (
`event_id` BIGINT NOT NULL,
`team_id` VARCHAR(50) NOT NULL,
FOREIGN KEY (`event_id`) REFERENCES `event` (`id`) ON DELETE CASCADE,
FOREIGN KEY (`team_id`) REFERENCES `team` (`name`) ON DELETE CASCADE
) CHARACTER SET utf8mb4 COMMENT='How participants relate';
""".strip(), # noqa
)
def get_app_connection(config, app) -> BaseDBAsyncClient:
"""
get connection name
:param config:
:param app:
:return:
"""
return Tortoise.get_connection(get_app_connection_name(config, app))
async def main():
await init_db()
# Generate the schema
conn = Tortoise.get_connection('default')
await conn.execute_query('CREATE EXTENSION IF NOT EXISTS citext')
await Tortoise.generate_schemas()
sql = 'ALTER TABLE "user" ALTER COLUMN username TYPE CITEXT'
await conn.execute_query(sql)
async def run():
print("SQLite:\n")
await Tortoise.init(db_url="sqlite://:memory:", modules={"models": ["__main__"]})
sql = get_schema_sql(Tortoise.get_connection("default"), safe=False)
print(sql)
print("\n\nMySQL:\n")
await Tortoise.init(db_url="mysql://root:@127.0.0.1:3306/", modules={"models": ["__main__"]})
sql = get_schema_sql(Tortoise.get_connection("default"), safe=False)
print(sql)
print("\n\nPostgreSQL:\n")
await Tortoise.init(
db_url="postgres://postgres:@127.0.0.1:5432/", modules={"models": ["__main__"]}
)
sql = get_schema_sql(Tortoise.get_connection("default"), safe=False)
print(sql)
def get_app_connection(config, app):
"""
get connection name
:param config:
:param app:
:return:
"""
return Tortoise.get_connection(get_app_connection_name(config, app))
async def test_simple_insert(self):
conn = Tortoise.get_connection("models")
await conn.execute_query("INSERT INTO author (name) VALUES ('Foo')")
self.assertEqual(
await conn.execute_query_dict("SELECT name FROM author"),
[{
"name": "Foo"
}])
def setUp(self) -> None:
client = Tortoise.get_connection("models")
if client.schema_generator is MySQLSchemaGenerator:
self.ddl = MysqlDDL(client)
elif client.schema_generator is SqliteSchemaGenerator:
self.ddl = SqliteDDL(client)
elif client.schema_generator is AsyncpgSchemaGenerator:
self.ddl = PostgresDDL(client)
def db(self):
from tortoise import Tortoise
if self.default_connection not in current_transaction_map:
return None
return (
current_transaction_map[self.default_connection].get()
or Tortoise.get_connection(self.default_connection)
)
def db(self) -> BaseDBAsyncClient:
from tortoise import Tortoise
if self.default_connection not in current_transaction_map:
raise ConfigurationError('No DB associated to model')
return (
current_transaction_map[self.default_connection].get()
or Tortoise.get_connection(self.default_connection) # type: ignore
)
async def test_in_transaction_rollback(self):
async with in_transaction() as conn:
await conn.execute_query("INSERT INTO author (name) VALUES ('Foo')"
)
await conn.rollback()
conn = Tortoise.get_connection("models")
self.assertEqual(
await conn.execute_query_dict("SELECT name FROM author"), [])
