neo4j-etl-lib 0.1.0__tar.gz → 0.1.1__tar.gz

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neo4j-etl-lib
-Version: 0.1.0
+Version: 0.1.1
 Summary: Building blocks for ETL pipelines.
 Keywords: etl,graph,database
 Author-email: Bert Radke <bert.radke@pm.me>

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.1.1}/src/etl_lib/__init__.py

@@ -1,4 +1,4 @@
 """
 Building blocks for ETL pipelines.
 """
-__version__ = "0.1.0"
+__version__ = "0.1.1"

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.1.1}/src/etl_lib/cli/run_tools.py

@@ -55,7 +55,7 @@ def __driver(ctx):
     database_name = ctx.obj["database_name"]
     neo4j_password = ctx.obj["neo4j_password"]
     return GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_password), database=database_name,
-                                notifications_min_severity="OFF", user_agent="ETL CLI 0.1")
+                                notifications_min_severity="OFF", user_agent="ETL CLI")
 
 
 @click.group()
@@ -165,25 +165,57 @@ def detail(ctx, run_id, details):
             __print_details(driver, run_id)
 
 
+# noinspection PyTypeChecker
 @cli.command()
-@click.option('--run-id', required=False, help='Run ID to delete')
-@click.option('--since', help='Delete runs since a specific date')
-@click.option('--older', help='Delete runs older than a specific date')
+@click.option('--run-id', required=False, type=str, help='Run IDs to delete, works with comma separated list')
+@click.option('--before', type=click.DateTime(formats=["%Y-%m-%d"]), help='Delete runs before a specific date in format YYYY-MM-DD')
+@click.option('--older', help='Delete runs older than x days', type=int)
 @click.pass_context
-def delete(ctx, run_id, since, older):
+def delete(ctx, run_id, before, older):
     """
-    Delete runs based on run ID, date, or age. One and only one of --run-id, --since, or --older must be provided.
+    Delete runs based on run ID, date, or age. One and only one of --run-id, --before, or --older must be provided.
     """
     # Ensure mutual exclusivity
-    options = [run_id, since, older]
+    options = [run_id, before, older]
     if sum(bool(opt) for opt in options) != 1:
-        print("You must specify exactly one of --run-id, --since, or --older.")
+        print("You must specify exactly one of --run-id, --before, or --older.")
         return
 
     if run_id:
-        print(f"Deleting run ID: {run_id}")
-    elif since:
-        print(f"Deleting runs since: {since}")
+        ids = run_id.split(',')
+        delete_runs(ctx, ids)
+    elif before:
+        print(f"Deleting runs before: {before}")
+        with __driver(ctx) as driver:
+            record= driver.execute_query(
+                """MATCH (r:ETLRun) WHERE date(r.startTime) < date($before) 
+                RETURN collect(r.uuid) AS ids
+                """,
+                result_transformer_=neo4j.Result.single,
+                before=before)
+            ids = record[0]
+            delete_runs(ctx, ids)
+
     elif older:
         print(f"Deleting runs older than: {older}")
-    # Implement delete logic here
+        with __driver(ctx) as driver:
+            record = driver.execute_query(
+                """MATCH (r:ETLRun) WHERE date(r.startTime) < (date() - duration({days: $days})) 
+                RETURN collect(r.uuid) AS ids
+                """,
+                result_transformer_=neo4j.Result.single,
+                days=older)
+            ids = record[0]
+            delete_runs(ctx, ids)
+
+
+def delete_runs(ctx, ids):
+    print(f"Deleting run IDs: {ids}")
+    with __driver(ctx) as driver:
+        records, _, _ = driver.execute_query(
+            """
+            MATCH (r:ETLRun)-[*]->(n) WHERE r.uuid IN $ids
+            DETACH DELETE n
+            DETACH DELETE r
+            """, ids=ids, routing_=neo4j.RoutingControl.WRITE)
+    print(f"Deleted run IDs: {ids}  successfully")

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.1.1}/src/etl_lib/core/ProgressReporter.py

@@ -3,7 +3,7 @@ from datetime import datetime
 
 from tabulate import tabulate
 
-from etl_lib.core.Task import Task, TaskGroup
+from etl_lib.core.Task import Task, TaskGroup, TaskReturn
 
 
 class ProgressReporter:
@@ -47,7 +47,7 @@ class ProgressReporter:
         self.logger.info(f"{'\t' * task.depth}starting {task.task_name()}")
         return task
 
-    def finished_task(self, task: Task, success: bool, summery: dict, error: str = None) -> Task:
+    def finished_task(self, task: Task, result: TaskReturn) -> Task:
         """
         Marks the task as finished.
 
@@ -55,23 +55,21 @@ class ProgressReporter:
 
         Args:
             task: Task to be marked as finished.
-            success: True if the task has successfully finished.
-            summery: statistics for this task (such as `nodes_created`)
-            error: If an exception occurred, the exception text should be provided here.
+            result: result of the task execution, such as status and summery information.
 
         Returns:
             Task to be marked as started.
         """
         task.end_time = datetime.now()
-        task.success = success
-        task.summery = summery
+        task.success = result.success
+        task.summery = result.summery
 
-        report = f"{'\t' * task.depth} finished {task.task_name()} in {task.end_time - task.start_time} with success: {success}"
-        if error is not None:
-            report += f", error: \n{error}"
+        report = f"{'\t' * task.depth} finished {task.task_name()} in {task.end_time - task.start_time} with success: {result.success}"
+        if result.error is not None:
+            report += f", error: \n{result.error}"
         else:
             # for the logger, remove entries with 0, but keep them in the original for reporting
-            cleaned_summery = {key: value for key, value in summery.items() if value != 0}
+            cleaned_summery = {key: value for key, value in result.summery.items() if value != 0}
             if len(cleaned_summery) > 0:
                 report += f"\n{tabulate([cleaned_summery], headers='keys', tablefmt='psql')}"
         self.logger.info(report)
@@ -168,9 +166,9 @@ class Neo4jProgressReporter(ProgressReporter):
                         start_time=task.start_time)
         return task
 
-    def finished_task(self, task: Task, success: bool, summery: dict, error: str = None) -> Task:
-        super().finished_task(task=task, success=success, summery=summery, error=error)
-        if success:
+    def finished_task(self, task: Task,  result: TaskReturn) -> Task:
+        super().finished_task(task=task, result=result)
+        if result.success:
             status = "success"
         else:
             status = "failure"
@@ -179,7 +177,7 @@ class Neo4jProgressReporter(ProgressReporter):
             MATCH (t:ETLTask {uuid:$id}) SET t.endTime = $end_time, t.status = $status, t.error = $error
             CREATE (s:ETLStats) SET s=$summery
             CREATE (t)-[:HAS_STATS]->(s)
-            """, id=task.uuid, end_time=task.end_time, summery=summery, status=status, error=error)
+            """, id=task.uuid, end_time=task.end_time, summery=result.summery, status=status, error=result.error)
         return task
 
     def __create_constraints(self):

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.1.1}/src/etl_lib/core/Task.py

@@ -46,7 +46,8 @@ class TaskReturn:
 
         # Combine success values and errors
         combined_success = self.success and other.success
-        combined_error = f"{self.error or ''} | {other.error or ''}".strip(" |")
+        combined_error = None if not (self.error or other.error) \
+            else f"{self.error or ''} | {other.error or ''}".strip(" |")
 
         return TaskReturn(
             success=combined_success, summery=merged_summery, error=combined_error
@@ -99,12 +100,7 @@ class Task:
         except Exception as e:
             result = TaskReturn(success=False, summery={}, error=str(e))
 
-        self.context.reporter.finished_task(
-            task=self,
-            success=result.success,
-            summery=result.summery,
-            error=result.error,
-        )
+        self.context.reporter.finished_task(task=self,result=result)
 
         return result
 

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.1.1}/src/etl_lib/data_source/CypherBatchSource.py

@@ -1,4 +1,6 @@
-from typing import Generator
+from typing import Generator, Callable, Optional
+
+from neo4j import Record
 
 from etl_lib.core.BatchProcessor import BatchResults, BatchProcessor
 from etl_lib.core.ETLContext import ETLContext
@@ -7,7 +9,14 @@ from etl_lib.core.Task import Task
 
 class CypherBatchSource(BatchProcessor):
 
-    def __init__(self, context: ETLContext, task: Task, query: str, **kwargs):
+    def __init__(
+            self,
+            context: ETLContext,
+            task: Task,
+            query: str,
+            record_transformer: Optional[Callable[[Record], dict]] = None,
+            **kwargs
+    ):
         """
        Constructs a new CypherBatchSource.
 
@@ -15,10 +24,12 @@ class CypherBatchSource(BatchProcessor):
            context: :class:`etl_lib.core.ETLContext.ETLContext` instance.
            task: :class:`etl_lib.core.Task.Task` instance owning this batchProcessor.
            query: Cypher query to execute.
+           record_transformer: Optional function to transform each record. See Neo4j API documentation on `result_transformer_`
            kwargs: Arguments passed as parameters with the query.
        """
         super().__init__(context, task)
         self.query = query
+        self.record_transformer = record_transformer
         self.kwargs = kwargs
 
     def __read_records(self, tx, batch_size):
@@ -26,7 +37,11 @@ class CypherBatchSource(BatchProcessor):
         result = tx.run(self.query, **self.kwargs)
 
         for record in result:
-            batch_.append(record.data())
+            data = record.data()
+            if self.record_transformer:
+                data = self.record_transformer(data)
+            batch_.append(data)
+
             if len(batch_) == batch_size:
                 yield batch_
                 batch_ = []

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.1.1}/src/etl_lib/task/ExecuteCypherTask.py

@@ -24,10 +24,10 @@ class ExecuteCypherTask(Task):
                 for query in self._query():
                     result = self.context.neo4j.query_database(session=session, query=query, **kwargs)
                     stats = merge_summery(stats, result.summery)
-                return TaskReturn(True, stats)
+                return TaskReturn(success=True, summery=stats)
             else:
                 result = self.context.neo4j.query_database(session=session, query=self._query(), **kwargs)
-                return TaskReturn(True, result.summery)
+                return TaskReturn(success=True, summery=result.summery)
 
     @abc.abstractmethod
     def _query(self) -> str | list[str]:

{neo4j_etl_lib-0.1.0 → neo4j_etl_lib-0.1.1}/src/etl_lib/test_utils/utils.py

@@ -125,7 +125,7 @@ class DummyReporter:
     def started_task(self, task: Task) -> Task:
         pass
 
-    def finished_task(self, task, success: bool, summery: dict, error: str = None) -> Task:
+    def finished_task(self, task, result) -> Task:
         pass
 
     def report_progress(self, task, batches: int, expected_batches: int, stats: dict) -> None: