Add immune_tables config option (#68)

Author: droserasprout

poetry.lock

@@ -23,7 +23,6 @@ packages = [
 [tool.poetry.dependencies]
 python = "^3.8"
 aiohttp = "^3.7.4"
-aiomysql = "^0.0.21"
 asyncpg = "0.23.0"
 datamodel-code-generator = "^0.11.1"
 "ruamel.yaml" = "^0.17.2"

pyproject.toml

@@ -52,29 +52,6 @@ def connection_string(self) -> str:
         return f'{self.kind}://{self.path}'
 
 
-@dataclass
-class MySQLDatabaseConfig:
-    """MySQL database connection config
-
-    :param host: Host
-    :param port: Port
-    :param user: User
-    :param password: Password
-    :param database: Database name
-    """
-
-    kind: Literal['mysql']
-    host: str
-    port: int
-    user: str
-    database: str
-    password: str = ''
-
-    @property
-    def connection_string(self) -> str:
-        return f'{self.kind}://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}'
-
-
 @dataclass
 class PostgresDatabaseConfig:
     """Postgres database connection config
@@ -85,6 +62,7 @@ class PostgresDatabaseConfig:
     :param password: Password
     :param database: Database name
     :param schema_name: Schema name
+    :param immune_tables: List of tables to preserve during reindexing
     """
 
     kind: Literal['postgres']
@@ -94,11 +72,18 @@ class PostgresDatabaseConfig:
     database: str
     schema_name: str = 'public'
     password: str = ''
+    immune_tables: Optional[List[str]] = None
 
     @property
     def connection_string(self) -> str:
         return f'{self.kind}://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}?schema={self.schema_name}'
 
+    @validator('immune_tables')
+    def valid_immune_tables(cls, v):
+        if v and 'dipdup_state' in v:
+            raise ConfigurationError('`dipdup_state` table can\'t be immune')
+        return v
+
 
 @dataclass
 class ContractConfig:
@@ -635,7 +620,6 @@ class DipDupConfig:
     :param templates: Mapping of template aliases and index templates
     :param database: Database config
     :param hasura: Hasura config
-    :param configuration: Dynamic configuration parameters
     """
 
     spec_version: str
@@ -644,7 +628,7 @@ class DipDupConfig:
     contracts: Dict[str, ContractConfig] = Field(default_factory=dict)
     indexes: Dict[str, IndexConfigT] = Field(default_factory=dict)
     templates: Optional[Dict[str, IndexConfigTemplateT]] = None
-    database: Union[SqliteDatabaseConfig, MySQLDatabaseConfig, PostgresDatabaseConfig] = SqliteDatabaseConfig(kind='sqlite')
+    database: Union[SqliteDatabaseConfig, PostgresDatabaseConfig] = SqliteDatabaseConfig(kind='sqlite')
     hasura: Optional[HasuraConfig] = None
 
     def __post_init_post_parse__(self):

src/dipdup/config.py

@@ -1,9 +1,14 @@
+import os
+import sys
 from typing import Any, Dict, Optional
 
-from dipdup.config import ContractConfig, DipDupConfig, StaticTemplateConfig
+from tortoise import Tortoise
+from tortoise.transactions import in_transaction
+
+from dipdup.config import ContractConfig, DipDupConfig, PostgresDatabaseConfig, StaticTemplateConfig
 from dipdup.datasources import DatasourceT
 from dipdup.exceptions import ConfigurationError
-from dipdup.utils import FormattedLogger, reindex, restart
+from dipdup.utils import FormattedLogger
 
 
 # TODO: Dataclasses are cool, everyone loves them. Resolve issue with pydantic in HandlerContext.
@@ -34,11 +39,36 @@ def reset(self) -> None:
     def updated(self) -> bool:
         return self._updated
 
+    async def restart(self) -> None:
+        """Restart preserving CLI arguments"""
+        # NOTE: Remove --reindex from arguments to avoid reindexing loop
+        if '--reindex' in sys.argv:
+            sys.argv.remove('--reindex')
+        os.execl(sys.executable, sys.executable, *sys.argv)
+
     async def reindex(self) -> None:
-        await reindex()
+        """Drop all tables or whole database and restart with the same CLI arguments"""
+        if isinstance(self.config.database, PostgresDatabaseConfig):
+            exclude_expression = ''
+            if self.config.database.immune_tables:
+                immune_tables = [f"'{t}'" for t in self.config.database.immune_tables]
+                exclude_expression = f' AND tablename NOT IN ({",".join(immune_tables)})'
 
-    async def restart(self) -> None:
-        await restart()
+            async with in_transaction() as conn:
+                await conn.execute_script(
+                    f'''
+                    DO $$ DECLARE
+                        r RECORD;
+                    BEGIN
+                        FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = current_schema(){exclude_expression}) LOOP
+                            EXECUTE 'DROP TABLE IF EXISTS ' || quote_ident(r.tablename) || ' CASCADE';
+                        END LOOP;
+                    END $$;
+                    '''
+                )
+        else:
+            await Tortoise._drop_databases()
+        await self.restart()
 
     def add_contract(self, name: str, address: str, typename: Optional[str] = None) -> None:
         if name in self.config.contracts:

src/dipdup/context.py

@@ -255,7 +255,7 @@ async def _initialize_database(self, reindex: bool = False) -> None:
         # TODO: Move higher
         if reindex:
             self._logger.warning('Started with `--reindex` argument, reindexing')
-            await utils.reindex()
+            await self._ctx.reindex()
 
         try:
             schema_state = await State.get_or_none(index_type=IndexType.schema, index_name=connection_name)
@@ -268,7 +268,7 @@ async def _initialize_database(self, reindex: bool = False) -> None:
             await schema_state.save()
         elif schema_state.hash != schema_hash:
             self._logger.warning('Schema hash mismatch, reindexing')
-            await utils.reindex()
+            await self._ctx.reindex()
 
         sql_path = join(self._config.package_path, 'sql')
         if not exists(sql_path):

src/dipdup/dipdup.py

@@ -19,7 +19,7 @@
 from dipdup.context import HandlerContext
 from dipdup.datasources.tzkt.datasource import BigMapFetcher, OperationFetcher, TzktDatasource
 from dipdup.models import BigMapAction, BigMapData, BigMapDiff, OperationData, Origination, State, TemporaryState, Transaction
-from dipdup.utils import FormattedLogger, in_global_transaction, reindex
+from dipdup.utils import FormattedLogger, in_global_transaction
 
 OperationGroup = namedtuple('OperationGroup', ('hash', 'counter'))
 
@@ -82,7 +82,7 @@ async def _initialize_index_state(self) -> None:
 
         elif state.hash != index_hash:
             self._logger.warning('Config hash mismatch, reindexing')
-            await reindex()
+            await self._ctx.reindex()
 
         self._state = state
 

src/dipdup/index.py

@@ -1,8 +1,6 @@
 import asyncio
 import logging
-import os
 import re
-import sys
 import time
 from contextlib import asynccontextmanager
 from logging import Logger
@@ -118,35 +116,6 @@ async def http_request(session: aiohttp.ClientSession, method: str, **kwargs):
         return await response.json()
 
 
-async def restart() -> None:
-    """Restart preserving CLI arguments"""
-    # NOTE: Remove --reindex from arguments to avoid reindexing loop
-    if '--reindex' in sys.argv:
-        sys.argv.remove('--reindex')
-    os.execl(sys.executable, sys.executable, *sys.argv)
-
-
-async def reindex() -> None:
-    """Drop all tables or whole database and restart with the same CLI arguments"""
-    if isinstance(Tortoise._connections['default'], AsyncpgDBClient):
-        async with in_transaction() as conn:
-            await conn.execute_script(
-                '''
-                DO $$ DECLARE
-                    r RECORD;
-                BEGIN
-                    FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = current_schema()) LOOP
-                        EXECUTE 'DROP TABLE IF EXISTS ' || quote_ident(r.tablename) || ' CASCADE';
-                    END LOOP;
-                END $$;
-                '''
-            )
-    else:
-        await Tortoise._drop_databases()
-    # NOTE: Tortoise can't recover after dropping database for some reason, restart.
-    await restart()
-
-
 class FormattedLogger(Logger):
     def __init__(
         self,