Datasets modularization pt.4 (#441)

### Feature or Bugfix - Refactoring ### Detail Refactoring of DatasetTable: Get rid of ElasticSearch connection for every request. Created a lazy way to establish connection. ### Relates #412 and #295 By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license. --------- Co-authored-by: dbalintx <132444646+dbalintx@users.noreply.github.com>
data-dot-all · May 4, 2023 · 3a9b3d8 · 3a9b3d8
1 parent 2a4c3d6
commit 3a9b3d8
Show file tree

Hide file tree

Showing 61 changed files with 371 additions and 379 deletions.
diff --git a/backend/api_handler.py b/backend/api_handler.py
@@ -15,7 +15,6 @@
 from dataall.core.context import set_context, dispose_context, RequestContext
 from dataall.db import init_permissions, get_engine, api, permissions
 from dataall.modules.loader import load_modules, ImportMode
-from dataall.searchproxy import connect
 
 logger = logging.getLogger()
 logger.setLevel(os.environ.get('LOG_LEVEL', 'INFO'))
@@ -30,7 +29,6 @@
 TYPE_DEFS = gql(SCHEMA.gql(with_directives=False))
 ENVNAME = os.getenv('envname', 'local')
 ENGINE = get_engine(envname=ENVNAME)
-ES = connect(envname=ENVNAME)
 Worker.queue = SqsQueue.send
 
 init_permissions(ENGINE)
@@ -99,7 +97,6 @@ def handler(event, context):
 
     log.info('Lambda Event %s', event)
     log.debug('Env name %s', ENVNAME)
-    log.debug('ElasticSearch %s', ES)
     log.debug('Engine %s', ENGINE.engine.url)
 
     if event['httpMethod'] == 'OPTIONS':
@@ -137,11 +134,10 @@ def handler(event, context):
             print(f'Error managing groups due to: {e}')
             groups = []
 
-        set_context(RequestContext(ENGINE, username, groups, ES))
+        set_context(RequestContext(ENGINE, username, groups))
 
         app_context = {
             'engine': ENGINE,
-            'es': ES,
             'username': username,
             'groups': groups,
             'schema': SCHEMA,

diff --git a/backend/dataall/api/Objects/Dashboard/resolvers.py b/backend/dataall/api/Objects/Dashboard/resolvers.py
@@ -311,7 +311,7 @@ def delete_dashboard(context: Context, source, dashboardUri: str = None):
             data=None,
             check_perm=True,
         )
-        indexers.delete_doc(es=context.es, doc_id=dashboardUri)
+        DashboardIndexer.delete_doc(doc_id=dashboardUri)
         return True
 
 

diff --git a/backend/dataall/api/Objects/Dataset/resolvers.py b/backend/dataall/api/Objects/Dataset/resolvers.py
@@ -328,9 +328,7 @@ def sync_tables(context: Context, source, datasetUri: str = None):
         DatasetTableIndexer.upsert_all(
             session=session, dataset_uri=dataset.datasetUri
         )
-        indexers.remove_deleted_tables(
-            session=session, es=context.es, datasetUri=dataset.datasetUri
-        )
+        DatasetTableIndexer.remove_all_deleted(session=session, dataset_uri=dataset.datasetUri)
         return Dataset.paginated_dataset_tables(
             session=session,
             username=context.username,
@@ -557,13 +555,13 @@ def delete_dataset(
 
         tables = [t.tableUri for t in Dataset.get_dataset_tables(session, datasetUri)]
         for uri in tables:
-            indexers.delete_doc(es=context.es, doc_id=uri)
+            DatasetIndexer.delete_doc(doc_id=uri)
 
         folders = [f.locationUri for f in DatasetLocationService.get_dataset_folders(session, datasetUri)]
         for uri in folders:
-            indexers.delete_doc(es=context.es, doc_id=uri)
+            DatasetIndexer.delete_doc(doc_id=uri)
 
-        indexers.delete_doc(es=context.es, doc_id=datasetUri)
+        DatasetIndexer.delete_doc(doc_id=datasetUri)
 
         Dataset.delete_dataset(
             session=session,

diff --git a/backend/dataall/api/Objects/Feed/registry.py b/backend/dataall/api/Objects/Feed/registry.py
@@ -38,5 +38,4 @@ def types(cls):
 
 FeedRegistry.register(FeedDefinition("Worksheet", models.Worksheet))
 FeedRegistry.register(FeedDefinition("DataPipeline", models.DataPipeline))
-FeedRegistry.register(FeedDefinition("DatasetTable", models.DatasetTable))
 FeedRegistry.register(FeedDefinition("Dashboard", models.Dashboard))
diff --git a/backend/dataall/api/Objects/Glossary/registry.py b/backend/dataall/api/Objects/Glossary/registry.py
@@ -1,7 +1,5 @@
-from dataclasses import dataclass, field
-from typing import Type, Dict, Optional, Protocol, Union, Callable, Any
-
-from opensearchpy import OpenSearch
+from dataclasses import dataclass
+from typing import Type, Dict, Optional, Protocol, Union
 
 from dataall.api import gql
 from dataall.api.gql.graphql_union_type import UnionTypeRegistry
@@ -56,7 +54,7 @@ def types(cls):
         return [gql.Ref(definition.object_type) for definition in cls._DEFINITIONS.values()]
 
     @classmethod
-    def reindex(cls, session, es: OpenSearch, target_type: str, target_uri: str):
+    def reindex(cls, session, target_type: str, target_uri: str):
         definition = cls._DEFINITIONS[target_type]
         if definition.reindexer:
             definition.reindexer.upsert(session, target_uri)

diff --git a/backend/dataall/api/Objects/Glossary/resolvers.py b/backend/dataall/api/Objects/Glossary/resolvers.py
@@ -458,7 +458,7 @@ def reindex(context, linkUri):
         if not link:
             return
 
-    GlossaryRegistry.reindex(session, context.es, link.targetType, link.targetUri)
+    GlossaryRegistry.reindex(session, link.targetType, link.targetUri)
 
 
 def _target_model(target_type: str):

diff --git a/backend/dataall/api/Objects/ShareObject/resolvers.py b/backend/dataall/api/Objects/ShareObject/resolvers.py
@@ -7,7 +7,7 @@
 from ....api.context import Context
 from ....aws.handlers.service_handlers import Worker
 from ....db import models
-from dataall.modules.datasets.db.models import DatasetStorageLocation
+from dataall.modules.datasets.db.models import DatasetStorageLocation, DatasetTable
 
 log = logging.getLogger(__name__)
 
@@ -265,7 +265,7 @@ def resolve_dataset(context: Context, source: models.ShareObject, **kwargs):
 
 
 def union_resolver(object, *_):
-    if isinstance(object, models.DatasetTable):
+    if isinstance(object, DatasetTable):
         return 'DatasetTable'
     elif isinstance(object, DatasetStorageLocation):
         return 'DatasetStorageLocation'

diff --git a/backend/dataall/api/Objects/Vote/resolvers.py b/backend/dataall/api/Objects/Vote/resolvers.py
@@ -28,11 +28,11 @@ def upvote(context: Context, source, input=None):
             data=input,
             check_perm=True,
         )
-        reindex(session, context.es, vote)
+        reindex(session, vote)
         return vote
 
 
-def reindex(session, es, vote):
+def reindex(session, vote):
     if vote.targetType == 'dataset':
         DatasetIndexer.upsert(session=session, dataset_uri=vote.targetUri)
     elif vote.targetType == 'dashboard':

diff --git a/backend/dataall/api/Objects/__init__.py b/backend/dataall/api/Objects/__init__.py
@@ -17,7 +17,6 @@
     DataPipeline,
     Environment,
     Activity,
-    DatasetTable,
     Dataset,
     Group,
     Principal,

diff --git a/backend/dataall/api/context.py b/backend/dataall/api/context.py
@@ -2,11 +2,9 @@ class Context:
     def __init__(
         self,
         engine=None,
-        es=None,
         username=None,
         groups=None,
     ):
         self.engine = engine
-        self.es = es
         self.username = username
         self.groups = groups
diff --git a/backend/dataall/aws/handlers/glue.py b/backend/dataall/aws/handlers/glue.py
@@ -6,6 +6,7 @@
 from .sts import SessionHelper
 from ... import db
 from ...db import models
+from dataall.modules.datasets.db.models import DatasetTable
 
 log = logging.getLogger('aws:glue')
 
@@ -524,7 +525,7 @@ def get_job_runs(engine, task: models.Task):
 
     @staticmethod
     def grant_principals_all_table_permissions(
-        table: models.DatasetTable, principals: [str], client=None
+        table: DatasetTable, principals: [str], client=None
     ):
         """
         Update the table permissions on Lake Formation

diff --git a/backend/dataall/aws/handlers/redshift.py b/backend/dataall/aws/handlers/redshift.py
@@ -11,6 +11,7 @@
 from ...db import models
 # TODO should be migrated in the redshift module
 from dataall.modules.datasets.services.dataset_table import DatasetTableService
+from dataall.modules.datasets.db.models import DatasetTable
 
 log = logging.getLogger(__name__)
 
@@ -448,7 +449,7 @@ def copy_data(engine, task: models.Task):
                 session, task.payload['datasetUri']
             )
 
-            table: models.DatasetTable = DatasetTableService.get_dataset_table_by_uri(
+            table: DatasetTable = DatasetTableService.get_dataset_table_by_uri(
                 session, task.payload['tableUri']
             )
 

diff --git a/backend/dataall/core/context.py b/backend/dataall/core/context.py
@@ -12,7 +12,6 @@
 
 from dataall.db.connection import Engine
 from threading import local
-import opensearchpy
 
 
 _request_storage = local()
@@ -24,7 +23,6 @@ class RequestContext:
     db_engine: Engine
     username: str
     groups: List[str]
-    es_engine: opensearchpy.OpenSearch
 
 
 def get_context() -> RequestContext:

diff --git a/backend/dataall/db/api/dataset.py b/backend/dataall/db/api/dataset.py
@@ -16,9 +16,10 @@
 from . import Organization
 from .. import models, api, exceptions, permissions, paginate
 from ..models.Enums import Language, ConfidentialityClassification
-from ...modules.datasets.db.dataset_repository import DatasetRepository
-from ...modules.datasets.services.dataset_location import DatasetLocationService
-from ...utils.naming_convention import (
+from dataall.modules.datasets.db.dataset_repository import DatasetRepository
+from dataall.modules.datasets.db.models import DatasetTable
+from dataall.modules.datasets.services.dataset_location import DatasetLocationService
+from dataall.utils.naming_convention import (
     NamingConventionService,
     NamingConventionPattern,
 )
@@ -266,21 +267,21 @@ def paginated_dataset_tables(
         session, username, groups, uri, data=None, check_perm=None
     ) -> dict:
         query = (
-            session.query(models.DatasetTable)
+            session.query(DatasetTable)
             .filter(
                 and_(
-                    models.DatasetTable.datasetUri == uri,
-                    models.DatasetTable.LastGlueTableStatus != 'Deleted',
+                    DatasetTable.datasetUri == uri,
+                    DatasetTable.LastGlueTableStatus != 'Deleted',
                 )
             )
-            .order_by(models.DatasetTable.created.desc())
+            .order_by(DatasetTable.created.desc())
         )
         if data and data.get('term'):
             query = query.filter(
                 or_(
                     *[
-                        models.DatasetTable.name.ilike('%' + data.get('term') + '%'),
-                        models.DatasetTable.GlueTableName.ilike(
+                        DatasetTable.name.ilike('%' + data.get('term') + '%'),
+                        DatasetTable.GlueTableName.ilike(
                             '%' + data.get('term') + '%'
                         ),
                     ]
@@ -379,7 +380,7 @@ def transfer_stewardship_to_new_stewards(session, dataset, new_stewards):
                 group=new_stewards,
                 permissions=permissions.DATASET_TABLE_READ,
                 resource_uri=tableUri,
-                resource_type=models.DatasetTable.__name__,
+                resource_type=DatasetTable.__name__,
             )
 
         dataset_shares = (
@@ -455,8 +456,8 @@ def update_glue_database_status(session, dataset_uri):
     def get_dataset_tables(session, dataset_uri):
         """return the dataset tables"""
         return (
-            session.query(models.DatasetTable)
-            .filter(models.DatasetTable.datasetUri == dataset_uri)
+            session.query(DatasetTable)
+            .filter(DatasetTable.datasetUri == dataset_uri)
             .all()
         )
 
@@ -585,10 +586,10 @@ def _delete_dataset_term_links(session, uri):
     @staticmethod
     def _delete_dataset_tables(session, dataset_uri) -> bool:
         tables = (
-            session.query(models.DatasetTable)
+            session.query(DatasetTable)
             .filter(
                 and_(
-                    models.DatasetTable.datasetUri == dataset_uri,
+                    DatasetTable.datasetUri == dataset_uri,
                 )
             )
             .all()
@@ -618,7 +619,7 @@ def get_dataset_by_bucket_name(session, bucket) -> [models.Dataset]:
     @staticmethod
     def count_dataset_tables(session, dataset_uri):
         return (
-            session.query(models.DatasetTable)
-            .filter(models.DatasetTable.datasetUri == dataset_uri)
+            session.query(DatasetTable)
+            .filter(DatasetTable.datasetUri == dataset_uri)
             .count()
         )
diff --git a/backend/dataall/db/api/redshift_cluster.py b/backend/dataall/db/api/redshift_cluster.py
@@ -4,11 +4,12 @@
 
 from .. import models, api, exceptions, paginate, permissions
 from . import has_resource_perm, ResourcePolicy, Environment, Dataset
-from ...utils.naming_convention import (
+from dataall.modules.datasets.db.models import DatasetTable
+from dataall.utils.naming_convention import (
     NamingConventionService,
     NamingConventionPattern,
 )
-from ...utils.slugify import slugify
+from dataall.utils.slugify import slugify
 
 log = logging.getLogger(__name__)
 
@@ -334,13 +335,13 @@ def list_available_cluster_tables(
         )
         created = (
             session.query(
-                models.DatasetTable.datasetUri.label('datasetUri'),
-                models.DatasetTable.tableUri.label('tableUri'),
+                DatasetTable.datasetUri.label('datasetUri'),
+                DatasetTable.tableUri.label('tableUri'),
                 models.RedshiftCluster.clusterUri.label('clusterUri'),
             )
             .join(
                 models.Dataset,
-                models.DatasetTable.datasetUri == models.Dataset.datasetUri,
+                DatasetTable.datasetUri == models.Dataset.datasetUri,
             )
             .filter(
                 and_(
@@ -354,19 +355,19 @@ def list_available_cluster_tables(
                 )
             )
             .group_by(
-                models.DatasetTable.datasetUri,
-                models.DatasetTable.tableUri,
+                DatasetTable.datasetUri,
+                DatasetTable.tableUri,
                 models.RedshiftCluster.clusterUri,
             )
         )
         all_group_tables_sub_query = shared.union(created).subquery(
             'all_group_tables_sub_query'
         )
         query = (
-            session.query(models.DatasetTable)
+            session.query(DatasetTable)
             .join(
                 all_group_tables_sub_query,
-                all_group_tables_sub_query.c.tableUri == models.DatasetTable.tableUri,
+                all_group_tables_sub_query.c.tableUri == DatasetTable.tableUri,
             )
             .filter(
                 models.RedshiftCluster.clusterUri == cluster.clusterUri,
@@ -541,18 +542,18 @@ def list_copy_enabled_tables(
         session, username, groups, uri, data=None, check_perm=True
     ) -> [models.RedshiftClusterDatasetTable]:
         q = (
-            session.query(models.DatasetTable)
+            session.query(DatasetTable)
             .join(
                 models.RedshiftClusterDatasetTable,
                 models.RedshiftClusterDatasetTable.tableUri
-                == models.DatasetTable.tableUri,
+                == DatasetTable.tableUri,
             )
             .filter(models.RedshiftClusterDatasetTable.clusterUri == uri)
         )
         if data.get('term'):
             term = data.get('term')
             q = q.filter(
-                models.DatasetTable.label.ilike('%' + term + '%'),
+                DatasetTable.label.ilike('%' + term + '%'),
             )
         return paginate(
             q, page=data.get('page', 1), page_size=data.get('pageSize', 20)