добавляет dag

dags/task_01/composites.py

@@ -0,0 +1,22 @@
+from task_01.target_db.repositories import TargetDBRepo
+from task_01.tasks import UserActionTransfer
+
+
+class Adapters:
+
+    def __init__(self, driver):
+        self.driver = driver
+
+    @property
+    def target_repo(self):
+        return TargetDBRepo(self.driver)
+
+
+class UserActionTransferTask:
+
+    def __init__(self, driver):
+        self.adapters = Adapters(driver)
+
+    @property
+    def run(self):
+        return UserActionTransfer(target_db=self.adapters.target_repo, )

dags/task_01/dto.py

@@ -0,0 +1,8 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class ConnectionInfo:
+    uri: str
+    username: str
+    password: str

dags/task_01/target_db/repositories.py

@@ -0,0 +1,41 @@
+import logging
+
+from neo4j import Driver
+from pandas import DataFrame
+
+BATCH_SIZE = 1000
+
+
+class TargetDBRepo:
+
+    def __init__(self, driver: Driver):
+        self.driver = driver
+        self.log = logging.getLogger(__name__)
+
+    def save_users(self, users: DataFrame) -> None:
+        query = """
+        UNWIND $rows AS row
+        CREATE (u:User {
+            user_id: row.user_id,
+            action: row.action,
+            timestamp: row.timestamp
+        })
+        """
+
+        with self.driver.session() as session:
+            for i in range(0, len(users), BATCH_SIZE):
+                batch = users.iloc[i:i + BATCH_SIZE]
+                records = batch.to_dict(orient="records")
+                session.run(query, {"rows": records})
+                self.log.info("rows saved %s", i + BATCH_SIZE)
+
+    def get_number_of_users(self) -> int:
+        with self.driver.session() as session:
+            result = session.run(
+                "MATCH (u:User) RETURN count(u) as user_count")
+            return result.single()["user_count"]
+
+    def check_connection(self):
+        with self.driver.session() as session:
+            result = session.run('RETURN "Connection successful" AS message')
+            return result.single()["message"]

dags/task_01/target_db/utills.py

@@ -0,0 +1,21 @@
+from airflow.hooks.base import BaseHook
+from neo4j import Driver, GraphDatabase
+from task_01.dto import ConnectionInfo
+
+
+def get_neo4j_connection(conn_id: str) -> ConnectionInfo:
+    conn = BaseHook.get_connection(conn_id)
+    uri = f'bolt://{conn.host}:{conn.port}'
+    return ConnectionInfo(
+        uri,
+        username=conn.login,
+        password=conn.password,
+    )
+
+
+def get_driver(conn_id: str) -> Driver:
+    connection_info = get_neo4j_connection(conn_id)
+    return GraphDatabase.driver(
+        uri=connection_info.uri,
+        auth=(connection_info.username, connection_info.password),
+    )

dags/task_01/tasks.py

@@ -0,0 +1,84 @@
+import logging
+import os
+import tempfile
+from datetime import datetime
+from random import randint
+
+import pandas as pd
+from airflow import AirflowException
+from airflow.models import Variable
+from task_01.target_db.repositories import TargetDBRepo
+
+log = logging.getLogger(__name__)
+
+
+class UserActionTransfer:
+
+    def __init__(
+        self,
+        target_db: TargetDBRepo,
+    ):
+        self.target_db = target_db
+        self.log = logging.getLogger(__name__)
+
+    def generate_sample_data(self):
+        """Генерация тестовых данных и сохранение в CSV файл"""
+        actions = self.get_fake_user_action_data()
+        csv_file_path = self.get_csv_file_path()
+        actions.to_csv(csv_file_path, index=False)
+
+        # Сохраняем путь к файлу в переменную DAG
+        Variable.set("user_action_data_path", csv_file_path)
+
+        logging.info("Sample data generated and saved to: %s", csv_file_path)
+        logging.info(f"Data preview: %s", actions.head())
+
+    def load_data_to_neo4j(self) -> None:
+        csv_file_path = Variable.get("user_action_data_path")
+
+        if not csv_file_path or not os.path.exists(csv_file_path):
+            raise AirflowException("CSV file not found: %s", csv_file_path)
+
+        # Чтение CSV файла
+        user_actions = pd.read_csv(csv_file_path)
+        logging.info("Loaded CSV data with %s rows", len(user_actions))
+
+        self.target_db.save_users(user_actions)
+        total_rows = self.target_db.get_number_of_users()
+        logging.info("Total rows: %s", total_rows)
+
+        # Очистка временного файла
+        os.remove(csv_file_path)
+        logging.info("Temporary CSV file cleaned up")
+
+    def check_neo4j_connection(self):
+        """Проверка соединения с БД"""
+        try:
+            result = self.target_db.check_connection()
+            log.info(f"Neo4j message: {result}")
+            log.info("Neo4j connection is healthy")
+        except Exception as e:
+            log.error(f"Neo4j connection failed: {e}")
+            raise
+        return result
+
+    @staticmethod
+    def get_csv_file_path() -> str:
+        temp_dir = tempfile.gettempdir()
+        return os.path.join(temp_dir, "user_action_data.csv")
+
+    @staticmethod
+    def get_fake_user_action_data() -> pd.DataFrame:
+        actions = ["login", "purchase", "view", "logout", "search"]
+        ids = list(range(1, 10001))
+        action = []
+        timestamp = []
+        for _ in ids:
+            action.append(actions[randint(0, len(actions) - 1)])
+            timestamp.append(datetime.now())
+        sample_data = {
+            "user_id": ids,
+            "action": action,
+            "timestamp": timestamp
+        }
+        return pd.DataFrame(sample_data)