Hasura config and table customisation improvements (#312)

* hasura config: add_source

* hasura: only customise tables from source
adding source with pg_add_source

* use_prepared_statements = True

* use postgres url instead of connection_parameters

* Changelog, rename config field, update tests

* Apply review patch

Co-authored-by: Lev Gorodetskiy <[email protected]>

Author: 852Kerfunkle

CHANGELOG.md

@@ -1,5 +1,15 @@
 # Changelog
 
+## [unreleased]
+
+### Added
+
+* config: Added `hasura.create_source` flag to create PostgreSQL source if missing.
+
+### Fixed
+
+* hasura: Do not apply table customizations to tables from other sources.
+
 ## 5.1.7 - 2022-06-15
 
 ### Fixed

src/dipdup/config.py

@@ -129,6 +129,16 @@ def connection_string(self) -> str:
             connection_string += f'&schema={self.schema_name}'
         return connection_string
 
+    @cached_property
+    def hasura_connection_parameters(self) -> Dict[str, Any]:
+        return {
+            'username': self.user,
+            'password': self.password,
+            'database': self.database,
+            'host': self.host,
+            'port': self.port,
+        }
+
     @validator('immune_tables')
     def _valid_immune_tables(cls, v) -> None:
         for table in v:
@@ -992,6 +1002,7 @@ class HasuraConfig:
 
     :param url: URL of the Hasura instance.
     :param admin_secret: Admin secret of the Hasura instance.
+    :param create_source: Whether source should be added to Hasura if missing.
     :param source: Hasura source for DipDup to configure, others will be left untouched.
     :param select_limit: Row limit for unauthenticated queries.
     :param allow_aggregations: Whether to allow aggregations in unauthenticated queries.
@@ -1002,6 +1013,7 @@ class HasuraConfig:
 
     url: str
     admin_secret: Optional[str] = None
+    create_source: bool = False
     source: str = 'default'
     select_limit: int = 100
     allow_aggregations: bool = True

src/dipdup/hasura.py

@@ -121,18 +121,23 @@ async def configure(self, force: bool = False) -> None:
             self._logger.info('Metadata is up to date, no action required')
             return
 
-        # NOTE: Find chosen source and overwrite its tables
+        # NOTE: Find chosen source and overwrite its tables, create if missing and allowed.
         source_name = self._hasura_config.source
-        for source in metadata['sources']:
-            if source['name'] == source_name:
-                source['tables'] = await self._generate_source_tables_metadata()
-                break
-        else:
-            raise HasuraError(f'Source `{source_name}` not found in metadata')
+
+        if (source := self._get_source(metadata, source_name)) is None:
+            if not self._hasura_config.create_source:
+                raise HasuraError(f'Source `{source_name}` not found in metadata. Set `create_source` flag to create it.')
+
+            await self._create_source()
+            metadata = await self._fetch_metadata()
+            if (source := self._get_source(metadata, source_name)) is None:
+                raise HasuraError(f'Source `{source_name}` not found in metadata after creation.')
+
+        source['tables'] = await self._generate_source_tables_metadata()
         await self._replace_metadata(metadata)
 
         # NOTE: Apply table customizations before generating queries
-        await self._apply_table_customization()
+        await self._apply_table_customization(source)
         metadata = await self._fetch_metadata()
 
         # NOTE: Generate and apply queries and REST endpoints
@@ -161,6 +166,13 @@ async def configure(self, force: bool = False) -> None:
 
         self._logger.info('Hasura instance has been configured')
 
+    def _get_source(self, metadata: Dict[str, Any], name: str) -> Optional[Dict[str, Any]]:
+        for source in metadata['sources']:
+            if source['name'] == name:
+                return source
+        else:
+            return None
+
     async def _hasura_request(self, endpoint: str, json: Dict[str, Any]) -> Dict[str, Any]:
         self._logger.debug('Sending `%s` request: %s', endpoint, dump_json(json))
         try:
@@ -203,13 +215,34 @@ async def _healthcheck(self) -> None:
 
         self._logger.info('Connected to Hasura %s', version)
 
+    async def _create_source(self) -> Dict[str, Any]:
+        self._logger.info(f'Adding source `{self._hasura_config.source}`')
+        return await self._hasura_request(
+            endpoint='metadata',
+            json={
+                'type': 'pg_add_source',
+                'args': {
+                    'name': self._hasura_config.source,
+                    'configuration': {
+                        'connection_info': {
+                            'database_url': {
+                                'connection_parameters': self._database_config.hasura_connection_parameters,
+                            },
+                            'use_prepared_statements': True,
+                        }
+                    },
+                    'replace_configuration': True,
+                },
+            },
+        )
+
     async def _fetch_metadata(self) -> Dict[str, Any]:
         self._logger.info('Fetching existing metadata')
         return await self._hasura_request(
             endpoint='metadata',
             json={
-                "type": "export_metadata",
-                "args": {},
+                'type': 'export_metadata',
+                'args': {},
             },
         )
 
@@ -400,17 +433,22 @@ async def _get_fields(self, name: str = 'query_root') -> List[Field]:
 
         return fields
 
-    async def _apply_table_customization(self) -> None:
+    async def _apply_table_customization(self, source: Dict[str, Any]) -> None:
         """Convert table and column names to camelCase.
 
         Based on https://github.com/m-rgba/hasura-snake-to-camel
         """
 
+        # NOTE: Build a set of table names for our source
+        table_names = {t['table']['name'] for t in source['tables']}
         tables = await self._get_fields()
 
         # TODO: Bulk request
         for table in tables:
-            custom_root_fields = self._format_custom_root_fields(table)
+            if table.root not in table_names:
+                continue
+
+            custom_root_fields = self._format_custom_root_fields(table.root)
             columns = await self._get_fields(table.root)
             custom_column_names = self._format_custom_column_names(columns)
             args: Dict[str, Any] = {
@@ -479,9 +517,7 @@ def _format_rest_endpoint(self, query_name: str) -> Dict[str, Any]:
             "comment": None,
         }
 
-    def _format_custom_root_fields(self, table: Field) -> Dict[str, Any]:
-        table_name = table.root
-
+    def _format_custom_root_fields(self, table_name: str) -> Dict[str, Any]:
         def _fmt(fmt: str) -> str:
             if self._hasura_config.camel_case:
                 return humps.camelize(fmt.format(table_name))

tests/integration_tests/test_hasura.py

@@ -55,7 +55,11 @@ async def test_configure_hasura(self) -> None:
             hasura_container._container.reload()
             hasura_ip = hasura_container._container.attrs['NetworkSettings']['IPAddress']
 
-            config.hasura = HasuraConfig(f'http://{hasura_ip}:8080')
+            config.hasura = HasuraConfig(
+                url=f'http://{hasura_ip}:8080',
+                source='new_source',
+                create_source=True,
+            )
             hasura_gateway = HasuraGateway('demo_hic_et_nunc', config.hasura, config.database)
             await stack.enter_async_context(hasura_gateway)