Index templates, generate only required types (#5)

* Add reference config, convert contracts block to dict`

* Index templates

* Generate required types only, cleanup schemas

* Add templates paragraph to readme

* Pass template_values to handlers

Author: droserasprout

README.md

@@ -182,6 +182,45 @@ $ # edit `secrets.env` file, change credentials
 $ docker-compose up dipdup
 ```
 
+### Index templates
+
+Sometimes you need to run multiple indexes with similar configs whose only difference is contract addresses. In this case you can use index templates like this:
+
+```yaml
+templates:
+  trades:
+    kind: operation
+    datasource: tzkt_staging
+    contract: < dex >
+    handlers:
+      - callback: on_fa12_token_to_tez
+        pattern:
+          - destination: < dex >
+            entrypoint: tokenToTezPayment
+          - destination: < token >
+            entrypoint: transfer
+      - callback: on_fa20_tez_to_token
+        pattern:
+          - destination: < dex >
+            entrypoint: tezToTokenPayment
+          - destination: < token >
+            entrypoint: transfer
+
+indexes:
+  trades_fa12:
+    template: trades
+    values:
+      dex: FA12_dex
+      token: FA12_token
+
+  trades_fa20:
+    template: trades
+    values:
+      dex: FA20_dex
+      token: FA20_token
+```
+
+
 ### Optional: configure Hasura GraphQL Engine
 
 `init` command generates Hasura metadata JSON in the package root. You can use `configure-graphql` command to apply it to the running GraphQL Engine instance:

src/dipdup/cli.py

@@ -149,6 +149,7 @@ async def init(ctx):
     await codegen.generate_types(config)
     await codegen.generate_handlers(config)
     await codegen.generate_hasura_metadata(config)
+    await codegen.cleanup(config)
 
 
 @cli.command(help='Configure Hasura GraphQL Engine')

src/dipdup/codegen.py

@@ -6,7 +6,8 @@
 from contextlib import suppress
 from os import mkdir
 from os.path import dirname, exists, join
-from typing import List
+from shutil import rmtree
+from typing import List, Tuple
 
 from jinja2 import Template
 from tortoise import Model, fields
@@ -51,15 +52,15 @@ async def fetch_schemas(config: DipDupConfig):
 
     for contract_name, contract in config.contracts.items():
         _logger.info('Fetching schemas for contract `%s`', contract_name)
-        address_schemas_path = join(schemas_path, contract.address)
+        address_schemas_path = join(schemas_path, contract)
         with suppress(FileExistsError):
             mkdir(address_schemas_path)
 
         parameter_schemas_path = join(address_schemas_path, 'parameter')
         with suppress(FileExistsError):
             mkdir(parameter_schemas_path)
 
-        address_schemas_json = await datasource.fetch_jsonschemas(contract.address)
+        address_schemas_json = await datasource.fetch_jsonschemas(contract)
         for entrypoint_json in address_schemas_json['entrypoints']:
             entrypoint = entrypoint_json['name']
             entrypoint_schema = entrypoint_json['parameterSchema']
@@ -79,6 +80,14 @@ async def generate_types(config: DipDupConfig):
         with open(join(types_path, '__init__.py'), 'w'):
             pass
 
+    _logger.info('Determining required types')
+    required_types: List[Tuple[str, str, str]] = []
+    for index_config in config.indexes.values():
+        if isinstance(index_config, OperationIndexConfig):
+            for handler_config in index_config.handlers:
+                for pattern_config in handler_config.pattern:
+                    required_types.append((pattern_config.destination, 'parameter', pattern_config.entrypoint))
+
     for root, dirs, files in os.walk(schemas_path):
         types_root = root.replace(schemas_path, types_path)
 
@@ -95,6 +104,13 @@ async def generate_types(config: DipDupConfig):
             entrypoint_name = file[:-5]
             entrypoint_name_titled = entrypoint_name.title().replace('_', '')
 
+            type_type = types_root.split('/')[-1]
+            contract = types_root.split('/')[-2]
+
+            key = (contract, type_type, entrypoint_name)
+            if key not in required_types:
+                continue
+
             input_path = join(root, file)
             output_path = join(types_root, f'{entrypoint_name}.py')
             _logger.info('Generating parameter type for entrypoint `%s`', entrypoint_name)
@@ -254,3 +270,9 @@ async def generate_hasura_metadata(config: DipDupConfig):
     metadata_path = join(config.package_path, 'hasura_metadata.json')
     with open(metadata_path, 'w') as file:
         json.dump(metadata, file, indent=4)
+
+
+async def cleanup(config: DipDupConfig):
+    _logger.info('Cleaning up')
+    schemas_path = join(config.package_path, 'schemas')
+    rmtree(schemas_path)

src/dipdup/config.py

@@ -29,11 +29,12 @@ class SqliteDatabaseConfig:
     :param path: Path to .sqlite3 file, leave default for in-memory database
     """
 
+    kind: Literal['sqlite']
     path: str = ':memory:'
 
     @property
     def connection_string(self):
-        return f'sqlite://{self.path}'
+        return f'{self.kind}://{self.path}'
 
 
 @dataclass
@@ -48,7 +49,7 @@ class DatabaseConfig:
     :param database: Schema name
     """
 
-    driver: str
+    kind: Union[Literal['postgres'], Literal['mysql']]
     host: str
     port: int
     user: str
@@ -57,7 +58,7 @@ class DatabaseConfig:
 
     @property
     def connection_string(self):
-        return f'{self.driver}://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}'
+        return f'{self.kind}://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}'
 
 
 @dataclass
@@ -144,6 +145,7 @@ class OperationIndexConfig:
     def __post_init_post_parse__(self):
         self._state: Optional[State] = None
         self._rollback_fn: Optional[Callable] = None
+        self._template_values: Dict[str, str] = None
 
     def hash(self) -> str:
         return hashlib.sha256(
@@ -173,6 +175,14 @@ def rollback_fn(self) -> Callable:
     def rollback_fn(self, value: Callable) -> None:
         self._rollback_fn = value
 
+    @property
+    def template_values(self) -> Optional[Dict[str, str]]:
+        return self._template_values
+
+    @template_values.setter
+    def template_values(self, value: Dict[str, str]) -> None:
+        self._template_values = value
+
 
 @dataclass
 class BigmapdiffHandlerPatternConfig:
@@ -208,14 +218,13 @@ class BlockIndexConfig:
 
 
 @dataclass
-class ContractConfig:
-    """Contract config
+class IndexTemplateConfig:
+    template: str
+    values: Dict[str, str]
 
-    :param network: Corresponding network alias, only for sanity checks
-    :param address: Contract address
-    """
 
-    address: str
+IndexConfigT = Union[OperationIndexConfig, BigmapdiffIndexConfig, BlockIndexConfig, IndexTemplateConfig]
+IndexConfigTemplateT = Union[OperationIndexConfig, BigmapdiffIndexConfig, BlockIndexConfig]
 
 
 @dataclass
@@ -232,21 +241,33 @@ class DipDupConfig:
 
     spec_version: str
     package: str
-    contracts: Dict[str, ContractConfig]
+    contracts: Dict[str, str]
     datasources: Dict[str, Union[TzktDatasourceConfig]]
-    indexes: Dict[str, Union[OperationIndexConfig, BigmapdiffIndexConfig, BlockIndexConfig]]
-    database: Union[SqliteDatabaseConfig, DatabaseConfig] = SqliteDatabaseConfig()
+    indexes: Dict[str, IndexConfigT]
+    templates: Optional[Dict[str, IndexConfigTemplateT]] = None
+    database: Union[SqliteDatabaseConfig, DatabaseConfig] = SqliteDatabaseConfig(kind='sqlite')
 
     def __post_init_post_parse__(self):
         self._logger = logging.getLogger(__name__)
+        for index_name, index_config in self.indexes.items():
+            if isinstance(index_config, IndexTemplateConfig):
+                template = self.templates[index_config.template]
+                raw_template = json.dumps(template, default=pydantic_encoder)
+                for key, value in index_config.values.items():
+                    value_regex = r'<[ ]*' + key + r'[ ]*>'
+                    raw_template = re.sub(value_regex, value, raw_template)
+                json_template = json.loads(raw_template)
+                self.indexes[index_name] = template.__class__(**json_template)
+                self.indexes[index_name].template_values = index_config.values
+
         for index_config in self.indexes.values():
             if isinstance(index_config, OperationIndexConfig):
                 if index_config is None:
                     continue
-                index_config.contract = self.contracts[index_config.contract].address
+                index_config.contract = self.contracts[index_config.contract]
                 for handler in index_config.handlers:
                     for pattern in handler.pattern:
-                        pattern.destination = self.contracts[pattern.destination].address
+                        pattern.destination = self.contracts[pattern.destination]
 
     @property
     def package_path(self) -> str:

src/dipdup/datasources/tzkt/cache.py

@@ -39,7 +39,7 @@ def match_operation(self, pattern_config: OperationHandlerPatternConfig, operati
 
     async def process(
         self,
-        callback: Callable[[OperationHandlerConfig, List[OperationData], List[OperationData]], Awaitable[None]],
+        callback: Callable[[OperationIndexConfig, OperationHandlerConfig, List[OperationData], List[OperationData]], Awaitable[None]],
     ) -> int:
         keys = list(self._operations.keys())
         self._logger.info('Matching %s operation groups', len(keys))
@@ -55,7 +55,7 @@ async def process(
 
                 if len(matched_operations) == len(handler_config.pattern):
                     self._logger.info('Handler `%s` matched! %s', handler_config.callback, key)
-                    await callback(handler_config, matched_operations, operations)
+                    await callback(self._index_config, handler_config, matched_operations, operations)
                     if key in self._operations:
                         del self._operations[key]
 

src/dipdup/datasources/tzkt/datasource.py

@@ -231,6 +231,7 @@ async def add_subscription(self, address: str, types: Optional[List[str]] = None
 
     async def on_operation_match(
         self,
+        index_config: OperationIndexConfig,
         handler_config: OperationHandlerConfig,
         matched_operations: List[OperationData],
         operations: List[OperationData],
@@ -247,7 +248,7 @@ async def on_operation_match(
             )
             args.append(context)
 
-        await handler_config.callback_fn(*args, operations)
+        await handler_config.callback_fn(*args, operations, index_config.template_values)
 
     @classmethod
     def convert_operation(cls, operation_json: Dict[str, Any]) -> OperationData:

src/dipdup/templates/handler.py.j2

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Dict, List, Optional
 from dipdup.models import HandlerContext, OperationData
 
 from {{ package }}.models import *
@@ -11,5 +11,6 @@ async def {{ handler }}(
     {{ pattern.entrypoint }}: HandlerContext[{{ pattern.entrypoint.title().replace('_', '') }}],
 {%- endfor %}
     operations: List[OperationData],
+    template_values: Optional[Dict[str, str]] = None,
 ) -> None:
     ...

src/dipdup_hic_et_nunc/dipdup-docker.yml

@@ -2,18 +2,16 @@ spec_version: 0.0.1
 package: dipdup_hic_et_nunc
 
 database:
-  driver: postgres
+  kind: postgres
   host: db
   port: 5432
   user: dipdup
   database: dipdup
   password: ${POSTGRES_PASSWORD:-changeme}
 
 contracts:
-  HEN_objkts:
-    address: ${HEN_OBJKTS:-KT1Hkg5qeNhfwpKW4fXvq7HGZB9z2EnmCCA9}
-  HEN_minter:
-    address: ${HEN_MINTER:-KT1RJ6PbjHpwc3M5rw5s2Nbmefwbuwbdxton}
+  HEN_objkts: ${HEN_OBJKTS:-KT1Hkg5qeNhfwpKW4fXvq7HGZB9z2EnmCCA9}
+  HEN_minter: ${HEN_MINTER:-KT1RJ6PbjHpwc3M5rw5s2Nbmefwbuwbdxton}
 
 datasources:
   tzkt_mainnet:

src/dipdup_hic_et_nunc/dipdup.yml

@@ -2,13 +2,12 @@ spec_version: 0.0.1
 package: dipdup_hic_et_nunc
 
 database:
+  kind: sqlite
   path: db.sqlite3
 
 contracts:
-  HEN_objkts:
-    address: ${HEN_OBJKTS:-KT1Hkg5qeNhfwpKW4fXvq7HGZB9z2EnmCCA9}
-  HEN_minter:
-    address: ${HEN_MINTER:-KT1RJ6PbjHpwc3M5rw5s2Nbmefwbuwbdxton}
+  HEN_objkts: ${HEN_OBJKTS:-KT1Hkg5qeNhfwpKW4fXvq7HGZB9z2EnmCCA9}
+  HEN_minter: ${HEN_MINTER:-KT1RJ6PbjHpwc3M5rw5s2Nbmefwbuwbdxton}
 
 datasources:
   tzkt_mainnet:

src/dipdup_hic_et_nunc/handlers/on_mint.py

@@ -1,14 +1,15 @@
-from typing import List
-
+from typing import Dict, List, Optional
 from dipdup.models import HandlerContext, OperationData
+
 from dipdup_hic_et_nunc.models import *
+
 from dipdup_hic_et_nunc.types.KT1Hkg5qeNhfwpKW4fXvq7HGZB9z2EnmCCA9.parameter.mint_OBJKT import MintObjkt
 from dipdup_hic_et_nunc.types.KT1RJ6PbjHpwc3M5rw5s2Nbmefwbuwbdxton.parameter.mint import Mint
 
-
 async def on_mint(
     mint_OBJKT: HandlerContext[MintObjkt],
     mint: HandlerContext[Mint],
     operations: List[OperationData],
+    template_values: Optional[Dict[str, str]] = None,
 ) -> None:
-    await Address.get_or_create(address=mint.parameter.address)
+    ...