modularization: Datasets modularization pt.2

backend/dataall/api/Objects/__init__.py

@@ -29,7 +29,6 @@
     Test,
     SagemakerStudio,
     RedshiftCluster,
-    DatasetProfiling,
     Glossary,
     AthenaQueryResult,
     Worksheet,

backend/dataall/aws/handlers/glue.py

@@ -522,90 +522,6 @@ def get_job_runs(engine, task: models.Task):
                 return []
             return response['JobRuns']
 
-    @staticmethod
-    @Worker.handler('glue.job.start_profiling_run')
-    def start_profiling_run(engine, task: models.Task):
-        with engine.scoped_session() as session:
-            profiling: models.DatasetProfilingRun = (
-                db.api.DatasetProfilingRun.get_profiling_run(
-                    session, profilingRunUri=task.targetUri
-                )
-            )
-            dataset: models.Dataset = session.query(models.Dataset).get(
-                profiling.datasetUri
-            )
-            run = Glue.run_job(
-                **{
-                    'accountid': dataset.AwsAccountId,
-                    'name': dataset.GlueProfilingJobName,
-                    'region': dataset.region,
-                    'arguments': (
-                        {'--table': profiling.GlueTableName}
-                        if profiling.GlueTableName
-                        else {}
-                    ),
-                }
-            )
-            db.api.DatasetProfilingRun.update_run(
-                session,
-                profilingRunUri=profiling.profilingRunUri,
-                GlueJobRunId=run['JobRunId'],
-            )
-            return run
-
-    @staticmethod
-    def run_job(**data):
-        accountid = data['accountid']
-        name = data['name']
-        try:
-            session = SessionHelper.remote_session(accountid=accountid)
-            client = session.client('glue', region_name=data.get('region', 'eu-west-1'))
-            response = client.start_job_run(
-                JobName=name, Arguments=data.get('arguments', {})
-            )
-            return response
-        except ClientError as e:
-            log.error(f'Failed to start profiling job {name} due to: {e}')
-            raise e
-
-    @staticmethod
-    @Worker.handler('glue.job.profiling_run_status')
-    def get_profiling_run(engine, task: models.Task):
-        with engine.scoped_session() as session:
-            profiling: models.DatasetProfilingRun = (
-                db.api.DatasetProfilingRun.get_profiling_run(
-                    session, profilingRunUri=task.targetUri
-                )
-            )
-            dataset: models.Dataset = session.query(models.Dataset).get(
-                profiling.datasetUri
-            )
-            glue_run = Glue.get_job_run(
-                **{
-                    'accountid': dataset.AwsAccountId,
-                    'name': dataset.GlueProfilingJobName,
-                    'region': dataset.region,
-                    'run_id': profiling.GlueJobRunId,
-                }
-            )
-            profiling.status = glue_run['JobRun']['JobRunState']
-            session.commit()
-            return profiling.status
-
-    @staticmethod
-    def get_job_run(**data):
-        accountid = data['accountid']
-        name = data['name']
-        run_id = data['run_id']
-        try:
-            session = SessionHelper.remote_session(accountid=accountid)
-            client = session.client('glue', region_name=data.get('region', 'eu-west-1'))
-            response = client.get_job_run(JobName=name, RunId=run_id)
-            return response
-        except ClientError as e:
-            log.error(f'Failed to get job run {run_id} due to: {e}')
-            raise e
-
     @staticmethod
     def grant_principals_all_table_permissions(
         table: models.DatasetTable, principals: [str], client=None

backend/dataall/db/api/__init__.py

@@ -13,7 +13,6 @@
 from .share_object import ShareObject, ShareObjectSM, ShareItemSM
 from .dataset import Dataset
 from .dataset_location import DatasetStorageLocation
-from .dataset_profiling_run import DatasetProfilingRun
 from .notification import Notification
 from .redshift_cluster import RedshiftCluster
 from .vpc import Vpc

backend/dataall/db/models/__init__.py

@@ -5,11 +5,9 @@
 from .DashboardShare import DashboardShare
 from .DashboardShare import DashboardShareStatus
 from .Dataset import Dataset
-from .DatasetProfilingRun import DatasetProfilingRun
 from .DatasetQualityRule import DatasetQualityRule
 from .DatasetStorageLocation import DatasetStorageLocation
 from .DatasetTable import DatasetTable
-from .DatasetTableProfilingJob import DatasetTableProfilingJob
 from .Environment import Environment
 from .EnvironmentGroup import EnvironmentGroup
 from .FeedMessage import FeedMessage

backend/dataall/modules/datasets/__init__.py

@@ -4,7 +4,7 @@
 
 from dataall.api.Objects.Feed.registry import FeedRegistry, FeedDefinition
 from dataall.api.Objects.Glossary.registry import GlossaryRegistry, GlossaryDefinition
-from dataall.modules.datasets.db.table_column_model import DatasetTableColumn
+from dataall.modules.datasets.db.models import DatasetTableColumn
 from dataall.modules.loader import ModuleInterface, ImportMode
 
 log = logging.getLogger(__name__)

backend/dataall/modules/datasets/api/__init__.py

@@ -1,6 +1,7 @@
 """The GraphQL schema of datasets and related functionality"""
 from dataall.modules.datasets.api import (
-    table_column
+    table_column,
+    profiling
 )
 
-__all__ = ["table_column"]
+__all__ = ["table_column", "profiling"]

backend/dataall/api/Objects/DatasetProfiling/__init__.py → backend/dataall/modules/datasets/api/profiling/__init__.py

@@ -1,4 +1,4 @@
-from . import (
+from dataall.modules.datasets.api.profiling import (
     input_types,
     mutations,
     queries,

backend/dataall/api/Objects/DatasetProfiling/input_types.py → backend/dataall/modules/datasets/api/profiling/input_types.py

@@ -1,4 +1,4 @@
-from ... import gql
+from dataall.api import gql
 
 StartDatasetProfilingRunInput = gql.InputType(
     name='StartDatasetProfilingRunInput',

backend/dataall/api/Objects/DatasetProfiling/mutations.py → backend/dataall/modules/datasets/api/profiling/mutations.py

@@ -1,5 +1,8 @@
-from ... import gql
-from .resolvers import *
+from dataall.api import gql
+from dataall.modules.datasets.api.profiling.resolvers import (
+    start_profiling_run,
+    update_profiling_run_results
+)
 
 startDatasetProfilingRun = gql.MutationField(
     name='startDatasetProfilingRun',

backend/dataall/api/Objects/DatasetProfiling/queries.py → backend/dataall/modules/datasets/api/profiling/queries.py

@@ -1,5 +1,10 @@
-from ... import gql
-from .resolvers import *
+from dataall.api import gql
+from dataall.modules.datasets.api.profiling.resolvers import (
+    get_profiling_run,
+    list_profiling_runs,
+    list_table_profiling_runs,
+    get_last_table_profiling_run
+)
 
 
 getDatasetProfilingRun = gql.QueryField(

backend/dataall/api/Objects/DatasetProfiling/resolvers.py → backend/dataall/modules/datasets/api/profiling/resolvers.py

@@ -1,17 +1,19 @@
 import json
 import logging
 
-from ....api.context import Context
-from ....aws.handlers.service_handlers import Worker
-from ....aws.handlers.sts import SessionHelper
-from ....db import api, permissions, models
-from ....db.api import ResourcePolicy
+from dataall.api.context import Context
+from dataall.aws.handlers.service_handlers import Worker
+from dataall.aws.handlers.sts import SessionHelper
+from dataall.db import api, permissions, models
+from dataall.db.api import ResourcePolicy
 from dataall.modules.datasets.services.dataset_table import DatasetTableService
+from dataall.modules.datasets.services.dataset_profiling_service import DatasetProfilingService
+from dataall.modules.datasets.db.models import DatasetProfilingRun
 
 log = logging.getLogger(__name__)
 
 
-def resolve_dataset(context, source: models.DatasetProfilingRun):
+def resolve_dataset(context, source: DatasetProfilingRun):
     if not source:
         return None
     with context.engine.scoped_session() as session:
@@ -32,7 +34,7 @@ def start_profiling_run(context: Context, source, input: dict = None):
         )
         dataset = api.Dataset.get_dataset_by_uri(session, input['datasetUri'])
 
-        run = api.DatasetProfilingRun.start_profiling(
+        run = DatasetProfilingService.start_profiling(
             session=session,
             datasetUri=dataset.datasetUri,
             tableUri=input.get('tableUri'),
@@ -49,7 +51,7 @@ def start_profiling_run(context: Context, source, input: dict = None):
     return run
 
 
-def get_profiling_run_status(context: Context, source: models.DatasetProfilingRun):
+def get_profiling_run_status(context: Context, source: DatasetProfilingRun):
     if not source:
         return None
     with context.engine.scoped_session() as session:
@@ -61,7 +63,7 @@ def get_profiling_run_status(context: Context, source: models.DatasetProfilingRu
     return source.status
 
 
-def get_profiling_results(context: Context, source: models.DatasetProfilingRun):
+def get_profiling_results(context: Context, source: DatasetProfilingRun):
     if not source or source.results == {}:
         return None
     else:
@@ -70,28 +72,28 @@ def get_profiling_results(context: Context, source: models.DatasetProfilingRun):
 
 def update_profiling_run_results(context: Context, source, profilingRunUri, results):
     with context.engine.scoped_session() as session:
-        run = api.DatasetProfilingRun.update_run(
+        run = DatasetProfilingService.update_run(
             session=session, profilingRunUri=profilingRunUri, results=results
         )
         return run
 
 
 def list_profiling_runs(context: Context, source, datasetUri=None):
     with context.engine.scoped_session() as session:
-        return api.DatasetProfilingRun.list_profiling_runs(session, datasetUri)
+        return DatasetProfilingService.list_profiling_runs(session, datasetUri)
 
 
 def get_profiling_run(context: Context, source, profilingRunUri=None):
     with context.engine.scoped_session() as session:
-        return api.DatasetProfilingRun.get_profiling_run(
+        return DatasetProfilingService.get_profiling_run(
             session=session, profilingRunUri=profilingRunUri
         )
 
 
 def get_last_table_profiling_run(context: Context, source, tableUri=None):
     with context.engine.scoped_session() as session:
-        run: models.DatasetProfilingRun = (
-            api.DatasetProfilingRun.get_table_last_profiling_run(
+        run: DatasetProfilingRun = (
+            DatasetProfilingService.get_table_last_profiling_run(
                 session=session, tableUri=tableUri
             )
         )
@@ -112,7 +114,7 @@ def get_last_table_profiling_run(context: Context, source, tableUri=None):
 
             if not run.results:
                 run_with_results = (
-                    api.DatasetProfilingRun.get_table_last_profiling_run_with_results(
+                    DatasetProfilingService.get_table_last_profiling_run_with_results(
                         session=session, tableUri=tableUri
                     )
                 )
@@ -143,6 +145,6 @@ def get_profiling_results_from_s3(environment, dataset, table, run):
 
 def list_table_profiling_runs(context: Context, source, tableUri=None):
     with context.engine.scoped_session() as session:
-        return api.DatasetProfilingRun.list_table_profiling_runs(
+        return DatasetProfilingService.list_table_profiling_runs(
             session=session, tableUri=tableUri, filter={}
         )

backend/dataall/api/Objects/DatasetProfiling/schema.py → backend/dataall/modules/datasets/api/profiling/schema.py

@@ -1,5 +1,5 @@
-from ... import gql
-from .resolvers import (
+from dataall.api import gql
+from dataall.modules.datasets.api.profiling.resolvers import (
     resolve_dataset,
     get_profiling_run_status,
     get_profiling_results,

backend/dataall/modules/datasets/api/table_column/resolvers.py

@@ -6,7 +6,7 @@
 from dataall.db import paginate, permissions, models
 from dataall.db.api import ResourcePolicy
 from dataall.modules.datasets.services.dataset_table import DatasetTableService
-from dataall.modules.datasets.db.table_column_model import DatasetTableColumn
+from dataall.modules.datasets.db.models import DatasetTableColumn
 
 
 def list_table_columns(

backend/dataall/modules/datasets/db/table_column_model.py → backend/dataall/modules/datasets/db/models.py

@@ -1,7 +1,6 @@
 from sqlalchemy import Column, String
-
-from dataall.db import Base
-from dataall.db import Resource, utils
+from sqlalchemy.dialects.postgresql import JSON
+from dataall.db import Base, Resource, utils
 
 
 class DatasetTableColumn(Resource, Base):
@@ -21,3 +20,19 @@ class DatasetTableColumn(Resource, Base):
 
     def uri(self):
         return self.columnUri
+
+
+class DatasetProfilingRun(Resource, Base):
+    __tablename__ = 'dataset_profiling_run'
+    profilingRunUri = Column(
+        String, primary_key=True, default=utils.uuid('profilingrun')
+    )
+    datasetUri = Column(String, nullable=False)
+    GlueJobName = Column(String)
+    GlueJobRunId = Column(String)
+    GlueTriggerSchedule = Column(String)
+    GlueTriggerName = Column(String)
+    GlueTableName = Column(String)
+    AwsAccountId = Column(String)
+    results = Column(JSON, default={})
+    status = Column(String, default='Created')

backend/dataall/modules/datasets/handlers/glue_column_handler.py

@@ -5,7 +5,7 @@
 from dataall.aws.handlers.service_handlers import Worker
 from dataall.modules.datasets.aws.glue_table_client import GlueTableClient
 from dataall.modules.datasets.aws.lf_table_client import LakeFormationTableClient
-from dataall.modules.datasets.db.table_column_model import DatasetTableColumn
+from dataall.modules.datasets.db.models import DatasetTableColumn
 from dataall.modules.datasets.services.dataset_table import DatasetTableService
 
 log = logging.getLogger(__name__)