Execute SQL scripts from on_restart/on_reindex directories (#70)

Author: droserasprout

src/dipdup/codegen.py

@@ -80,6 +80,17 @@ async def create_package(self) -> None:
             with open(join(handlers_path, '__init__.py'), 'w'):
                 pass
 
+        self._logger.info('Creating `%s/sql` directory', self._config.package)
+        sql_path = join(self._config.package_path, 'sql')
+        with suppress(FileExistsError):
+            mkdir(sql_path)
+        sql_on_restart_path = join(sql_path, 'on_restart')
+        with suppress(FileExistsError):
+            mkdir(sql_on_restart_path)
+        sql_on_reindex_path = join(sql_path, 'on_reindex')
+        with suppress(FileExistsError):
+            mkdir(sql_on_reindex_path)
+
     async def fetch_schemas(self) -> None:
         """Fetch JSONSchemas for all contracts used in config"""
         self._logger.info('Creating `schemas` package')

src/dipdup/dipdup.py

@@ -9,7 +9,7 @@
 from genericpath import exists
 from tortoise import Tortoise
 from tortoise.exceptions import OperationalError
-from tortoise.transactions import in_transaction
+from tortoise.transactions import get_connection
 from tortoise.utils import get_schema_sql
 
 import dipdup.utils as utils
@@ -30,7 +30,7 @@
 from dipdup.datasources import DatasourceT
 from dipdup.datasources.bcd.datasource import BcdDatasource
 from dipdup.datasources.tzkt.datasource import TzktDatasource
-from dipdup.exceptions import HandlerImportError
+from dipdup.exceptions import ConfigurationError, HandlerImportError
 from dipdup.hasura import configure_hasura
 from dipdup.index import BigMapIndex, HandlerContext, Index, OperationIndex
 from dipdup.models import BigMapData, IndexType, OperationData, State
@@ -264,27 +264,39 @@ async def _initialize_database(self, reindex: bool = False) -> None:
 
         if schema_state is None:
             await Tortoise.generate_schemas()
+            await self._execute_sql_scripts(reindex=True)
+
             schema_state = State(index_type=IndexType.schema, index_name=connection_name, hash=schema_hash)
             await schema_state.save()
         elif schema_state.hash != schema_hash:
             self._logger.warning('Schema hash mismatch, reindexing')
             await self._ctx.reindex()
 
+        await self._execute_sql_scripts(reindex=False)
+
+    async def _execute_sql_scripts(self, reindex: bool) -> None:
+        """Execute SQL included with project"""
         sql_path = join(self._config.package_path, 'sql')
         if not exists(sql_path):
             return
+        if any(map(lambda p: p not in ('on_reindex', 'on_restart'), listdir(sql_path))):
+            raise ConfigurationError(
+                f'SQL scripts must be placed either to `{self._config.package}/sql/on_restart` or to `{self._config.package}/sql/on_reindex` directory'
+            )
         if not isinstance(self._config.database, PostgresDatabaseConfig):
-            self._logger.warning('Injecting raw SQL supported on PostgreSQL only')
+            self._logger.warning('Execution of user SQL scripts is supported on PostgreSQL only, skipping')
             return
 
-        for filename in listdir(sql_path):
+        sql_path = join(sql_path, 'on_reindex' if reindex else 'on_restart')
+        if not exists(sql_path):
+            return
+        self._logger.info('Executing SQL scripts from `%s`', sql_path)
+        for filename in sorted(listdir(sql_path)):
             if not filename.endswith('.sql'):
                 continue
 
             with open(join(sql_path, filename)) as file:
                 sql = file.read()
 
-            self._logger.info('Applying raw SQL from `%s`', filename)
-
-            async with in_transaction() as conn:
-                await conn.execute_query(sql)
+            self._logger.info('Executing `%s`', filename)
+            await get_connection(None).execute_script(sql)