data-syncmaster 0.1.1__py3-none-any.whl

syncmaster/schemas/v1/transfers/strategy.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from syncmaster.schemas.v1.transfer_types import FULL_TYPE, INCREMENTAL_TYPE
+
+
+class FullStrategy(BaseModel):
+    type: FULL_TYPE
+
+
+class IncrementalStrategy(BaseModel):
+    type: INCREMENTAL_TYPE

syncmaster/schemas/v1/types.py

@@ -0,0 +1,5 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+from pydantic import constr
+
+NameConstr = constr(min_length=1)

syncmaster/schemas/v1/users.py

@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+from pydantic import BaseModel, constr
+
+from syncmaster.db.models import GroupMemberRole
+from syncmaster.db.utils import Pagination
+from syncmaster.schemas.v1.page import MetaPageSchema, PageSchema
+
+
+class UpdateUserSchema(BaseModel):
+    username: constr(pattern=r"^[_a-z0-9]+$")  # noqa: F722
+
+
+class ReadGroupMember(BaseModel):
+    id: int
+    username: str
+    role: GroupMemberRole
+
+    class Config:
+        from_attributes = True
+
+
+class ReadUserSchema(BaseModel):
+    id: int
+    username: str
+    is_superuser: bool
+
+    class Config:
+        from_attributes = True
+
+
+class FullUserSchema(ReadGroupMember):
+    is_active: bool
+
+    class Config:
+        from_attributes = True
+
+
+class UserPageSchemaAsGroupMember(PageSchema):
+    items: list[ReadGroupMember]
+
+    @classmethod
+    def from_pagination(cls, pagination: Pagination):
+        return cls(
+            meta=MetaPageSchema(
+                page=pagination.page,
+                pages=pagination.pages,
+                page_size=pagination.page_size,
+                total=pagination.total,
+                has_next=pagination.has_next,
+                has_previous=pagination.has_previous,
+                next_page=pagination.next_page,
+                previous_page=pagination.previous_page,
+            ),
+            items=[
+                ReadGroupMember(
+                    id=user.id,
+                    username=user.username,
+                    role=role,
+                )
+                for user, role in pagination.items
+            ],
+        )
+
+
+class UserPageSchema(PageSchema):
+    items: list[ReadUserSchema]
+
+    @classmethod
+    def from_pagination(cls, pagination: Pagination):
+        return cls(
+            meta=MetaPageSchema(
+                page=pagination.page,
+                pages=pagination.pages,
+                page_size=pagination.page_size,
+                total=pagination.total,
+                has_next=pagination.has_next,
+                has_previous=pagination.has_previous,
+                next_page=pagination.next_page,
+                previous_page=pagination.previous_page,
+            ),
+            items=pagination.items,
+        )

syncmaster/worker/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0

syncmaster/worker/base.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+from celery import Task
+from sqlalchemy import create_engine
+
+from syncmaster.config import Settings
+
+
+class WorkerTask(Task):
+    def __init__(self) -> None:
+        self.settings = Settings()
+        self.engine = create_engine(
+            url=self.settings.build_db_connection_uri(driver="psycopg2"),
+        )

syncmaster/worker/config.py

@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+from celery import Celery
+
+from syncmaster.config import Settings
+from syncmaster.worker.base import WorkerTask
+
+settings = Settings()
+
+celery = Celery(
+    __name__,
+    broker=settings.build_rabbit_connection_uri(),
+    backend="db+" + settings.build_db_connection_uri(driver="psycopg2"),
+    task_cls=WorkerTask,
+    imports=[
+        "syncmaster.worker.transfer",
+    ],
+)

syncmaster/worker/controller.py

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+import logging
+from typing import Any
+
+from syncmaster.config import Settings
+from syncmaster.db.models import Connection, Transfer
+from syncmaster.dto.connections import (
+    HDFSConnectionDTO,
+    HiveConnectionDTO,
+    OracleConnectionDTO,
+    PostgresConnectionDTO,
+    S3ConnectionDTO,
+)
+from syncmaster.dto.transfers import (
+    HDFSTransferDTO,
+    HiveTransferDTO,
+    OracleTransferDTO,
+    PostgresTransferDTO,
+    S3TransferDTO,
+)
+from syncmaster.exceptions.connection import ConnectionTypeNotRecognizedError
+from syncmaster.worker.handlers.base import Handler
+from syncmaster.worker.handlers.file.hdfs import HDFSHandler
+from syncmaster.worker.handlers.file.s3 import S3Handler
+from syncmaster.worker.handlers.hive import HiveHandler
+from syncmaster.worker.handlers.oracle import OracleHandler
+from syncmaster.worker.handlers.postgres import PostgresHandler
+
+logger = logging.getLogger(__name__)
+
+
+connection_handler_proxy = {
+    "hive": (
+        HiveHandler,
+        HiveConnectionDTO,
+        HiveTransferDTO,
+    ),
+    "oracle": (
+        OracleHandler,
+        OracleConnectionDTO,
+        OracleTransferDTO,
+    ),
+    "postgres": (
+        PostgresHandler,
+        PostgresConnectionDTO,
+        PostgresTransferDTO,
+    ),
+    "s3": (
+        S3Handler,
+        S3ConnectionDTO,
+        S3TransferDTO,
+    ),
+    "hdfs": (
+        HDFSHandler,
+        HDFSConnectionDTO,
+        HDFSTransferDTO,
+    ),
+}
+
+
+class TransferController:
+    source_handler: Handler
+    target_handler: Handler
+
+    def __init__(
+        self,
+        transfer: Transfer,
+        source_connection: Connection,
+        source_auth_data: dict,
+        target_connection: Connection,
+        target_auth_data: dict,
+        settings: Settings,
+    ):
+        self.source_handler = self.get_handler(
+            connection_data=source_connection.data,
+            transfer_params=transfer.source_params,
+            connection_auth_data=source_auth_data,
+        )
+        self.target_handler = self.get_handler(
+            connection_data=target_connection.data,
+            transfer_params=transfer.target_params,
+            connection_auth_data=target_auth_data,
+        )
+        spark = settings.CREATE_SPARK_SESSION_FUNCTION(
+            settings,
+            target=self.target_handler.connection_dto,
+            source=self.source_handler.connection_dto,
+        )
+
+        self.source_handler.set_spark(spark)
+        self.target_handler.set_spark(spark)
+        logger.info("source connection = %s", self.source_handler)
+        logger.info("target connection = %s", self.target_handler)
+
+    def start_transfer(self) -> None:
+        self.source_handler.init_connection()
+        self.source_handler.init_reader()
+
+        self.target_handler.init_connection()
+        self.target_handler.init_writer()
+        logger.info("Source and target were initialized")
+
+        df = self.target_handler.normalize_column_name(self.source_handler.read())
+        logger.info("Data has been read")
+
+        self.target_handler.write(df)
+        logger.info("Data has been inserted")
+
+    def get_handler(
+        self,
+        connection_data: dict[str, Any],
+        connection_auth_data: dict,
+        transfer_params: dict[str, Any],
+    ) -> Handler:
+        connection_data.update(connection_auth_data)
+        handler_type = connection_data["type"]
+
+        if connection_handler_proxy.get(handler_type, None) is None:
+            raise ConnectionTypeNotRecognizedError
+
+        handler, connection_dto, transfer_dto = connection_handler_proxy[handler_type]
+
+        return handler(
+            connection_dto=connection_dto(**connection_data),
+            transfer_dto=transfer_dto(**transfer_params),
+        )

syncmaster/worker/handlers/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0

syncmaster/worker/handlers/base.py

@@ -0,0 +1,49 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+from abc import ABC
+
+from onetl.db import DBReader, DBWriter
+from pyspark.sql import SparkSession
+from pyspark.sql.dataframe import DataFrame
+
+from syncmaster.dto.connections import ConnectionDTO
+from syncmaster.dto.transfers import TransferDTO
+
+
+class Handler(ABC):
+    def __init__(
+        self,
+        connection_dto: ConnectionDTO,
+        transfer_dto: TransferDTO,
+        spark: SparkSession | None = None,
+    ) -> None:
+        self.spark = spark
+        self.reader: DBReader | None = None
+        self.writer: DBWriter | None = None
+        self.connection_dto = connection_dto
+        self.transfer_dto = transfer_dto
+
+    def init_connection(self): ...
+
+    def set_spark(self, spark: SparkSession):
+        self.spark = spark
+
+    def init_reader(self):
+        if self.connection_dto is None:
+            raise ValueError("At first you need to initialize connection. Run `init_connection")
+
+    def init_writer(self):
+        if self.connection_dto is None:
+            raise ValueError("At first you need to initialize connection. Run `init_connection")
+
+    def read(self) -> DataFrame:
+        if self.reader is None:
+            raise ValueError("Reader is not initialized")
+        return self.reader.run()
+
+    def write(self, df: DataFrame) -> None:
+        if self.writer is None:
+            raise ValueError("Writer is not initialized")
+        return self.writer.run(df=df)
+
+    def normalize_column_name(self, df: DataFrame) -> DataFrame: ...

syncmaster/worker/handlers/file/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0

syncmaster/worker/handlers/file/base.py

@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+import json
+
+from onetl.base.base_file_df_connection import BaseFileDFConnection
+from onetl.file import FileDFReader, FileDFWriter
+from onetl.file.format import CSV, JSON, JSONLine
+from pyspark.sql.dataframe import DataFrame
+from pyspark.sql.types import StructType
+
+from syncmaster.dto.connections import ConnectionDTO
+from syncmaster.dto.transfers import TransferDTO
+from syncmaster.worker.handlers.base import Handler
+
+
+class FileHandler(Handler):
+    connection: BaseFileDFConnection
+    connection_dto: ConnectionDTO
+    transfer_dto: TransferDTO
+
+    def init_connection(self): ...
+
+    def init_reader(self):
+        super().init_reader()
+
+        self.reader = FileDFReader(
+            connection=self.connection,
+            format=self._get_format(),
+            source_path=self.transfer_dto.directory_path,
+            df_schema=StructType.fromJson(json.loads(self.transfer_dto.df_schema)),
+            options=self.transfer_dto.options,
+        )
+
+    def init_writer(self):
+        super().init_writer()
+
+        self.writer = FileDFWriter(
+            connection=self.connection,
+            format=self._get_format(),
+            target_path=self.transfer_dto.directory_path,
+            options=self.transfer_dto.options,
+        )
+
+    def normalize_column_name(self, df: DataFrame) -> DataFrame:
+        return df
+
+    def _get_format(self):
+        file_type = self.transfer_dto.file_format["type"]
+        if file_type == "csv":
+            return CSV.parse_obj(self.transfer_dto.file_format)
+        elif file_type == "jsonline":
+            return JSONLine.parse_obj(self.transfer_dto.file_format)
+        elif file_type == "json":
+            return JSON.parse_obj(self.transfer_dto.file_format)
+        else:
+            raise ValueError("Unknown file type")

syncmaster/worker/handlers/file/hdfs.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+
+from onetl.connection import SparkHDFS
+
+from syncmaster.worker.handlers.file.base import FileHandler
+
+
+class HDFSHandler(FileHandler):
+    def init_connection(self):
+        self.connection = SparkHDFS(
+            cluster=self.connection_dto.cluster,
+            spark=self.spark,
+        ).check()

syncmaster/worker/handlers/file/s3.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+from onetl.connection import SparkS3
+
+from syncmaster.worker.handlers.file.base import FileHandler
+
+
+class S3Handler(FileHandler):
+    def init_connection(self):
+        self.connection = SparkS3(
+            host=self.connection_dto.host,
+            port=self.connection_dto.port,
+            access_key=self.connection_dto.access_key,
+            secret_key=self.connection_dto.secret_key,
+            bucket=self.connection_dto.bucket,
+            protocol=self.connection_dto.protocol,
+            region=self.connection_dto.region,
+            extra=self.connection_dto.additional_params,
+            spark=self.spark,
+        ).check()

syncmaster/worker/handlers/hive.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+from onetl.connection import Hive
+from onetl.db import DBReader, DBWriter
+from pyspark.sql.dataframe import DataFrame
+
+from syncmaster.dto.connections import HiveConnectionDTO
+from syncmaster.dto.transfers import HiveTransferDTO
+from syncmaster.worker.handlers.base import Handler
+
+
+class HiveHandler(Handler):
+    connection: Hive
+    connection_dto: HiveConnectionDTO
+    transfer_dto: HiveTransferDTO
+
+    def init_connection(self):
+        self.connection = Hive(
+            cluster=self.connection_dto.cluster,
+            spark=self.spark,
+        ).check()
+
+    def init_reader(self):
+        super().init_reader()
+        self.spark.catalog.refreshTable(self.transfer_dto.table_name)
+        self.reader = DBReader(
+            connection=self.connection,
+            table=self.transfer_dto.table_name,
+        )
+
+    def init_writer(self):
+        super().init_writer()
+        self.writer = DBWriter(
+            connection=self.connection,
+            table=self.transfer_dto.table_name,
+        )
+
+    def normalize_column_name(self, df: DataFrame) -> DataFrame:
+        for column_name in df.columns:
+            df = df.withColumnRenamed(column_name, column_name.lower())
+        return df

syncmaster/worker/handlers/oracle.py

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+from onetl.connection import Oracle
+from onetl.db import DBReader, DBWriter
+from pyspark.sql.dataframe import DataFrame
+
+from syncmaster.dto.connections import OracleConnectionDTO
+from syncmaster.dto.transfers import OracleTransferDTO
+from syncmaster.worker.handlers.base import Handler
+
+
+class OracleHandler(Handler):
+    connection: Oracle
+    connection_dto: OracleConnectionDTO
+    transfer_dto: OracleTransferDTO
+
+    def init_connection(self):
+        self.connection = Oracle(
+            host=self.connection_dto.host,
+            port=self.connection_dto.port,
+            user=self.connection_dto.user,
+            password=self.connection_dto.password,
+            sid=self.connection_dto.sid,
+            service_name=self.connection_dto.service_name,
+            extra=self.connection_dto.additional_params,
+            spark=self.spark,
+        ).check()
+
+    def init_reader(self):
+        super().init_reader()
+        df = self.connection.get_df_schema(self.transfer_dto.table_name)
+        self.reader = DBReader(
+            connection=self.connection,
+            table=self.transfer_dto.table_name,
+            columns=[f'"{f}"' for f in df.fieldNames()],
+        )
+
+    def init_writer(self):
+        super().init_writer()
+        self.writer = DBWriter(
+            connection=self.connection,
+            table=self.transfer_dto.table_name,
+        )
+
+    def normalize_column_name(self, df: DataFrame) -> DataFrame:
+        for column_name in df.columns:
+            df = df.withColumnRenamed(column_name, column_name.upper())
+        return df

syncmaster/worker/handlers/postgres.py

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+from onetl.connection import Postgres
+from onetl.db import DBReader, DBWriter
+from pyspark.sql.dataframe import DataFrame
+
+from syncmaster.dto.connections import PostgresConnectionDTO
+from syncmaster.dto.transfers import PostgresTransferDTO
+from syncmaster.worker.handlers.base import Handler
+
+
+class PostgresHandler(Handler):
+    connection: Postgres
+    connection_dto: PostgresConnectionDTO
+    transfer_dto: PostgresTransferDTO
+
+    def init_connection(self):
+        self.connection = Postgres(
+            host=self.connection_dto.host,
+            user=self.connection_dto.user,
+            password=self.connection_dto.password,
+            port=self.connection_dto.port,
+            database=self.connection_dto.database_name,
+            extra=self.connection_dto.additional_params,
+            spark=self.spark,
+        ).check()
+
+    def init_reader(self):
+        super().init_reader()
+        df = self.connection.get_df_schema(self.transfer_dto.table_name)
+        self.reader = DBReader(
+            connection=self.connection,
+            table=self.transfer_dto.table_name,
+            columns=[f'"{f}"' for f in df.fieldNames()],
+        )
+
+    def init_writer(self):
+        super().init_writer()
+        self.writer = DBWriter(
+            connection=self.connection,
+            table=self.transfer_dto.table_name,
+        )
+
+    def normalize_column_name(self, df: DataFrame) -> DataFrame:
+        for column_name in df.columns:
+            df = df.withColumnRenamed(column_name, column_name.lower())
+        return df

syncmaster/worker/spark.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS (Mobile Telesystems)
+# SPDX-License-Identifier: Apache-2.0
+import logging
+
+import pyspark
+from onetl.connection import Oracle, Postgres, SparkS3
+from pyspark.sql import SparkSession
+
+from syncmaster.config import Settings
+from syncmaster.dto.connections import ConnectionDTO
+
+log = logging.getLogger(__name__)
+
+
+def get_worker_spark_session(
+    settings: Settings,  # used in test spark session definition
+    source: ConnectionDTO,
+    target: ConnectionDTO,
+) -> SparkSession:
+    """Through the source and target parameters you can get credentials for authorization at the source"""
+    spark_builder = SparkSession.builder.appName("celery_worker")
+
+    for k, v in get_spark_session_conf(source, target).items():
+        spark_builder = spark_builder.config(k, v)
+
+    if source.type == "hive" or target.type == "hive":  # type: ignore
+        log.debug("Enabling Hive support")
+        spark_builder = spark_builder.enableHiveSupport()
+
+    return spark_builder.getOrCreate()
+
+
+def get_packages(db_type: str) -> list[str]:
+    if db_type == "postgres":
+        return Postgres.get_packages()
+    if db_type == "oracle":
+        return Oracle.get_packages()
+    if db_type == "s3":
+        spark_version = pyspark.__version__
+        return SparkS3.get_packages(spark_version=spark_version)
+
+    # If the database type does not require downloading .jar packages
+    return []
+
+
+def get_excluded_packages(db_type: str):
+    if db_type == "s3":
+        return [
+            "com.google.cloud.bigdataoss:gcs-connector",
+            "org.apache.hadoop:hadoop-aliyun",
+            "org.apache.hadoop:hadoop-azure-datalake",
+            "org.apache.hadoop:hadoop-azure",
+        ]
+    return []
+
+
+def get_spark_session_conf(
+    source: ConnectionDTO,
+    target: ConnectionDTO,
+) -> dict:
+    maven_packages: list[str] = []
+    excluded_packages: list[str] = []
+
+    for db_type in source, target:
+        maven_packages.extend(get_packages(db_type=db_type.type))  # type: ignore
+        excluded_packages.extend(get_excluded_packages(db_type=db_type.type))  # type: ignore
+
+    log.debug("Passing Maven packages: %s", maven_packages)
+
+    config = {
+        "spark.jars.packages": ",".join(maven_packages),
+        "spark.sql.pyspark.jvmStacktrace.enabled": "true",
+    }
+
+    if excluded_packages:
+        config["spark.jars.excludes"] = ",".join(excluded_packages)
+
+    if source.type == "s3":  # type: ignore
+        config.update(
+            {
+                "spark.hadoop.fs.s3a.committer.magic.enabled": "true",
+                "spark.hadoop.fs.s3a.committer.name": "magic",
+                "spark.hadoop.mapreduce.outputcommitter.factory.scheme.s3a": (
+                    "org.apache.hadoop.fs.s3a.commit.S3ACommitterFactory"
+                ),
+                "spark.sql.parquet.output.committer.class": (
+                    "org.apache.spark.internal.io.cloud.BindingParquetOutputCommitter"
+                ),
+                "spark.sql.sources.commitProtocolClass": "org.apache.spark.internal.io.cloud.PathOutputCommitProtocol",
+            }
+        )
+
+    return config