Skip to content

Commit 6e0193c

Browse files
droserasproutdroserasprout
authored
Move main loop code from cli module (#37)
1 parent 624bc35 commit 6e0193c

File tree

3 files changed

+94
-77
lines changed

3 files changed

+94
-77
lines changed

‎src/dipdup/cli.py‎

Lines changed: 6 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,15 @@
11
import asyncio
2-
import hashlib
32
import logging
43
import os
5-
import sys
64
from dataclasses import dataclass
75
from functools import wraps
86
from os.path import dirname, join
9-
from typing import Dict
107

118
import click
12-
from tortoise import Tortoise
13-
from tortoise.exceptions import OperationalError
14-
from tortoise.utils import get_schema_sql
159

16-
import dipdup.codegen as codegen
1710
from dipdup import __version__
18-
from dipdup.config import DipDupConfig, IndexTemplateConfig, LoggingConfig, PostgresDatabaseConfig, TzktDatasourceConfig
19-
from dipdup.datasources.tzkt.datasource import TzktDatasource
20-
from dipdup.hasura import configure_hasura
21-
from dipdup.models import IndexType, State
22-
from dipdup.utils import reindex, tortoise_wrapper
11+
from dipdup.config import DipDupConfig, LoggingConfig
12+
from dipdup.dipdup import DipDup
2313

2414
_logger = logging.getLogger(__name__)
2515

@@ -68,69 +58,14 @@ async def cli(ctx, config: str, logging_config: str):
6858
@click_async
6959
async def run(ctx) -> None:
7060
config: DipDupConfig = ctx.obj.config
71-
72-
url = config.database.connection_string
73-
models = f'{config.package}.models'
74-
async with tortoise_wrapper(url, models):
75-
_logger.info('Initializing database')
76-
77-
connection_name, connection = next(iter(Tortoise._connections.items()))
78-
schema_sql = get_schema_sql(connection, False)
79-
80-
if isinstance(config.database, PostgresDatabaseConfig) and config.database.schema_name:
81-
await Tortoise._connections['default'].execute_script("CREATE SCHEMA IF NOT EXISTS {}".format(config.database.schema_name))
82-
await Tortoise._connections['default'].execute_script("SET search_path TO {}".format(config.database.schema_name))
83-
84-
# NOTE: Column order could differ in two generated schemas for the same models, drop commas and sort strings to eliminate this
85-
processed_schema_sql = '\n'.join(sorted(schema_sql.replace(',', '').split('\n'))).encode()
86-
schema_hash = hashlib.sha256(processed_schema_sql).hexdigest()
87-
88-
try:
89-
schema_state = await State.get_or_none(index_type=IndexType.schema, index_name=connection_name)
90-
except OperationalError:
91-
schema_state = None
92-
93-
if schema_state is None:
94-
await Tortoise.generate_schemas()
95-
schema_state = State(index_type=IndexType.schema, index_name=connection_name, hash=schema_hash)
96-
await schema_state.save()
97-
elif schema_state.hash != schema_hash:
98-
_logger.warning('Schema hash mismatch, reindexing')
99-
await reindex()
100-
101-
await config.initialize()
102-
103-
_logger.info('Fetching indexer state for dapp `%s`', config.package)
104-
datasources: Dict[TzktDatasourceConfig, TzktDatasource] = {}
105-
106-
for index_name, index_config in config.indexes.items():
107-
assert not isinstance(index_config, IndexTemplateConfig)
108-
_logger.info('Processing index `%s`', index_name)
109-
if isinstance(index_config.datasource, TzktDatasourceConfig):
110-
if index_config.tzkt_config not in datasources:
111-
datasources[index_config.tzkt_config] = TzktDatasource(index_config.tzkt_config.url)
112-
datasources[index_config.tzkt_config].add_index(index_config)
113-
else:
114-
raise NotImplementedError(f'Datasource `{index_config.datasource}` is not supported')
115-
116-
_logger.info('Starting datasources')
117-
run_tasks = [asyncio.create_task(d.start()) for d in datasources.values()]
118-
119-
if config.hasura:
120-
hasura_task = asyncio.create_task(configure_hasura(config))
121-
run_tasks.append(hasura_task)
122-
123-
await asyncio.gather(*run_tasks)
61+
dipdup = DipDup(config)
62+
await dipdup.run()
12463

12564

12665
@cli.command(help='Initialize new dipdap')
12766
@click.pass_context
12867
@click_async
12968
async def init(ctx):
13069
config: DipDupConfig = ctx.obj.config
131-
132-
await codegen.create_package(config)
133-
await codegen.fetch_schemas(config)
134-
await codegen.generate_types(config)
135-
await codegen.generate_handlers(config)
136-
await codegen.cleanup(config)
70+
dipdup = DipDup(config)
71+
await dipdup.init()

‎src/dipdup/dipdup.py‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import asyncio
2+
import hashlib
3+
import logging
4+
from typing import Dict
5+
6+
from tortoise import Tortoise
7+
from tortoise.exceptions import OperationalError
8+
from tortoise.utils import get_schema_sql
9+
10+
import dipdup.codegen as codegen
11+
from dipdup import __version__
12+
from dipdup.config import DipDupConfig, IndexTemplateConfig, PostgresDatabaseConfig, TzktDatasourceConfig
13+
from dipdup.datasources.tzkt.datasource import TzktDatasource
14+
from dipdup.hasura import configure_hasura
15+
from dipdup.models import IndexType, State
16+
from dipdup.utils import reindex, tortoise_wrapper
17+
18+
19+
class DipDup:
20+
def __init__(self, config: DipDupConfig) -> None:
21+
self._logger = logging.getLogger(__name__)
22+
self._config = config
23+
24+
async def init(self) -> None:
25+
await codegen.create_package(self._config)
26+
await codegen.fetch_schemas(self._config)
27+
await codegen.generate_types(self._config)
28+
await codegen.generate_handlers(self._config)
29+
await codegen.cleanup(self._config)
30+
31+
async def run(self) -> None:
32+
url = self._config.database.connection_string
33+
models = f'{self._config.package}.models'
34+
async with tortoise_wrapper(url, models):
35+
await self.initialize_database()
36+
37+
await self._config.initialize()
38+
39+
datasources: Dict[TzktDatasourceConfig, TzktDatasource] = {}
40+
41+
for index_name, index_config in self._config.indexes.items():
42+
assert not isinstance(index_config, IndexTemplateConfig)
43+
self._logger.info('Processing index `%s`', index_name)
44+
if isinstance(index_config.datasource, TzktDatasourceConfig):
45+
if index_config.tzkt_config not in datasources:
46+
datasources[index_config.tzkt_config] = TzktDatasource(index_config.tzkt_config.url)
47+
datasources[index_config.tzkt_config].add_index(index_config)
48+
else:
49+
raise NotImplementedError(f'Datasource `{index_config.datasource}` is not supported')
50+
51+
self._logger.info('Starting datasources')
52+
run_tasks = [asyncio.create_task(d.start()) for d in datasources.values()]
53+
54+
if self._config.hasura:
55+
hasura_task = asyncio.create_task(configure_hasura(self._config))
56+
run_tasks.append(hasura_task)
57+
58+
await asyncio.gather(*run_tasks)
59+
60+
async def initialize_database(self) -> None:
61+
self._logger.info('Initializing database')
62+
63+
if isinstance(self._config.database, PostgresDatabaseConfig) and self._config.database.schema_name:
64+
await Tortoise._connections['default'].execute_script(f"CREATE SCHEMA IF NOT EXISTS {self._config.database.schema_name}")
65+
await Tortoise._connections['default'].execute_script(f"SET search_path TO {self._config.database.schema_name}")
66+
67+
connection_name, connection = next(iter(Tortoise._connections.items()))
68+
schema_sql = get_schema_sql(connection, False)
69+
70+
# NOTE: Column order could differ in two generated schemas for the same models, drop commas and sort strings to eliminate this
71+
processed_schema_sql = '\n'.join(sorted(schema_sql.replace(',', '').split('\n'))).encode()
72+
schema_hash = hashlib.sha256(processed_schema_sql).hexdigest()
73+
74+
try:
75+
schema_state = await State.get_or_none(index_type=IndexType.schema, index_name=connection_name)
76+
except OperationalError:
77+
schema_state = None
78+
79+
if schema_state is None:
80+
await Tortoise.generate_schemas()
81+
schema_state = State(index_type=IndexType.schema, index_name=connection_name, hash=schema_hash)
82+
await schema_state.save()
83+
elif schema_state.hash != schema_hash:
84+
self._logger.warning('Schema hash mismatch, reindexing')
85+
await reindex()

‎tests/integration_tests/test_codegen.py‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@
55
from shutil import rmtree
66
from unittest import IsolatedAsyncioTestCase
77

8-
from dipdup import codegen
98
from dipdup.config import DipDupConfig
9+
from dipdup.dipdup import DipDup
1010

1111

1212
# NOTE: https://gist.github.com/breeze1990/0253cb96ce04c00cb7a67feb2221e95e
@@ -40,11 +40,8 @@ async def test_codegen(self):
4040
del sys.modules[config.package]
4141

4242
try:
43-
await codegen.create_package(config)
44-
await codegen.fetch_schemas(config)
45-
await codegen.generate_types(config)
46-
await codegen.generate_handlers(config)
47-
await codegen.cleanup(config)
43+
dipdup = DipDup(config)
44+
await dipdup.init()
4845

4946
import_submodules(config.package)
5047

0 commit comments

Comments
 (0)