recurvedata-lib 0.1.487__py2.py3-none-any.whl

recurvedata/pigeon/loader/csv_to_microsoft_fabric.py

@@ -0,0 +1,242 @@
+from typing import TYPE_CHECKING, Any, List, Optional, Union
+
+from recurvedata.pigeon import const
+from recurvedata.pigeon.loader.base import BaseLoader, CSVToDBAPIMixin
+from recurvedata.pigeon.utils import ensure_query_list, ensure_str_list, fs
+
+if TYPE_CHECKING:
+    from recurvedata.pigeon.connector.microsoft_fabric import MsFabricConnector
+
+allowed_modes = (const.LOAD_OVERWRITE, const.LOAD_MERGE, const.LOAD_APPEND)
+
+
+class CSVToMsFabricLoader(BaseLoader, CSVToDBAPIMixin):
+    """Loader for Microsoft Fabric that supports bulk loading data using COPY command.
+
+    This loader provides Microsoft Fabric specific data loading capabilities.
+    It uses the COPY command for efficient data loading and supports various
+    loading modes (OVERWRITE, MERGE, APPEND).
+    """
+
+    def __init__(
+        self,
+        table: str,
+        filename: str,
+        connector: "MsFabricConnector",
+        schema: Optional[str] = None,
+        create_table_ddl: Optional[str] = None,
+        mode: str = const.LOAD_MERGE,
+        primary_keys: Optional[Union[str, List[str]]] = None,
+        columns: Optional[Union[str, List[str]]] = None,
+        compress: bool = True,
+        delete_file: bool = True,
+        dedup: bool = False,
+        dedup_uniq_keys: Optional[Union[str, List[str]]] = None,
+        dedup_orderby: Optional[Union[str, List[str]]] = None,
+        pre_queries: Optional[Union[str, List[str]]] = None,
+        post_queries: Optional[Union[str, List[str]]] = None,
+        lineterminator: Optional[str] = "0x0D0A",
+        *args: Any,
+        **kwargs: Any,
+    ):
+        """Initialize the loader.
+
+        Args:
+            table: Target table name
+            filename: Source file path
+            connector: MsFabricConnector instance
+            schema: Schema name
+            create_table_ddl: SQL to create target table
+            mode: Loading mode (OVERWRITE/MERGE/APPEND)
+            primary_keys: Primary key columns for MERGE mode
+            columns: Column list to load
+            compress: Whether to compress data before loading
+            delete_file: Whether to delete source file after loading
+            dedup: Whether to deduplicate data
+            dedup_uniq_keys: Columns for deduplication
+            dedup_orderby: Order by clause for deduplication
+            pre_queries: Queries to run before loading
+            post_queries: Queries to run after loading
+        """
+        if "." in table:
+            self.schema, self.table = table.split(".")
+        else:
+            self.schema = schema or "dbo"
+            self.table = table
+
+        self.connector = connector
+        self.filename = filename
+        self.create_table_ddl = create_table_ddl
+        self.compress = compress
+        self.delete_file = delete_file
+
+        if mode not in allowed_modes:
+            raise ValueError(f"mode should be one of ({allowed_modes})")
+
+        self.mode = mode
+        self.primary_keys = ensure_str_list(primary_keys)
+        self.columns = ensure_str_list(columns)
+
+        # dedup stuff
+        self.dedup = dedup
+        self.dedup_uniq_keys = ensure_str_list(dedup_uniq_keys)
+        self.dedup_orderby = dedup_orderby
+        if self.dedup and not self.dedup_uniq_keys:
+            raise ValueError("dedup_uniq_keys should not be empty if dedup is true")
+
+        self.pre_queries = ensure_query_list(pre_queries) or []
+        self.post_queries = ensure_query_list(post_queries) or []
+        self.lineterminator = lineterminator
+
+        super().__init__()
+
+    @property
+    def staging_table(self) -> str:
+        return f"{self.table}_staging"
+
+    @property
+    def full_staging_table_name(self) -> str:
+        return f"{self.schema}.{self.staging_table}"
+
+    @property
+    def full_table_name(self) -> str:
+        return f"{self.schema}.{self.table}"
+
+    @property
+    def quoted_full_staging_table(self) -> str:
+        return self.connector.quote_identifier(self.full_staging_table_name)
+
+    @property
+    def quoted_full_table_name(self) -> str:
+        return self.connector.quote_identifier(self.full_table_name)
+
+    def execute_impl(self) -> None:
+        """Execute the data loading process."""
+        if fs.is_file_empty(self.filename):
+            self.logger.error("file not exists or has no content. %s", self.filename)
+            fs.remove_files_safely(fs.schema_filename(self.filename))
+            return
+
+        self._prepare_target_table()
+        self._prepare_staging_table()
+        self._load_to_staging()
+        self._merge_into_target_table()
+
+        # do cleaning things
+        if self.delete_file:
+            self.logger.info("delete local file %s", self.filename)
+            fs.remove_files_safely(self.filename)
+            fs.remove_files_safely(fs.schema_filename(self.filename))
+
+    def _prepare_staging_table(self) -> None:
+        """Prepare the staging table for data loading."""
+        schema, table = self.full_staging_table_name.split(".")
+        drop = self._make_drop_table_query(schema, table)
+        ddl = f"SELECT TOP 0 * INTO {self.quoted_full_staging_table} FROM {self.quoted_full_table_name}"
+        self.connector.execute([drop, ddl])
+
+    def _make_drop_table_query(self, schema: str, table: str) -> str:
+        """Generate SQL to drop a table if it exists."""
+        if "." in table:
+            schema, table = table.split(".")
+        if not schema:
+            schema = self.schema
+        full_table = f"{schema}.{table}"
+        query = f"""
+            IF EXISTS (
+              SELECT * FROM sys.tables
+              WHERE schema_name(schema_id) = '{schema}' AND name = '{table}'
+            )
+            DROP TABLE {self.connector.quote_identifier(full_table)}
+        """
+        return query
+
+    def _load_to_staging(self) -> None:
+        """Load data into staging table using COPY command."""
+        self.logger.info(f"load {self.filename} into staging table {self.full_staging_table_name}")
+        self.connector.load_csv_bulk(
+            table=self.full_staging_table_name,
+            filename=self.filename,
+            columns=self.columns,
+            compress=self.compress,
+            lineterminator=self.lineterminator,
+        )
+
+        if self.dedup:
+            dedup_query = self._construct_dedup_query()
+            self.connector.execute(dedup_query, autocommit=False, commit_on_close=True)
+
+    def _construct_dedup_query(self) -> str:
+        """Construct query for deduplication."""
+        partitions_cols = []
+        for col in self.dedup_uniq_keys:
+            partitions_cols.append(self.connector.quote_identifier(col))
+        partition_by = ", ".join(partitions_cols)
+        columns = " ,".join(self.connector.get_columns(schema=self.schema, table=self.staging_table))
+        tmp_table = f"{self.full_staging_table_name}_tmp"
+        quoted_tmp_table = self.connector.quote_identifier(tmp_table)
+        quoted_bak_table = self.connector.quote_identifier(f"{self.staging_table}_bak")
+
+        queries = f"""
+            {self._make_drop_table_query(self.schema, tmp_table)};
+
+            CREATE TABLE {quoted_tmp_table} AS
+            SELECT {', '.join(self.connector.quote_identifier(x) for x in columns)}
+            FROM (
+              SELECT *, ROW_NUMBER() OVER (PARTITION BY {partition_by} ORDER BY {self.dedup_orderby}) rn
+              FROM {self.quoted_full_staging_table}
+            ) AS t
+            WHERE rn = 1;
+
+            RENAME OBJECT {self.quoted_full_staging_table} TO {quoted_bak_table};
+            RENAME OBJECT {quoted_tmp_table} TO {self.staging_table};
+            DROP TABLE {quoted_bak_table};
+        """
+        return queries
+
+    def _merge_into_target_table(self) -> None:
+        """Merge data from staging table into target table."""
+        target = self.quoted_full_table_name
+        staging = self.quoted_full_staging_table
+
+        append_sql = f"INSERT INTO {target} SELECT * FROM {staging}"
+        if self.mode == const.LOAD_OVERWRITE:
+            queries = [f"TRUNCATE TABLE {target}", append_sql]
+        elif self.mode == const.LOAD_MERGE:
+            queries = self._ingest_by_merging()
+        else:
+            # APPEND
+            queries = [append_sql]
+
+        queries.append(f"DROP TABLE {staging}")
+
+        queries = self.pre_queries + queries + self.post_queries
+        self.logger.info("running Microsoft Fabric queries...")
+        self.connector.execute(queries, autocommit=True, commit_on_close=True)
+        self.logger.info("done.")
+
+    def _ingest_by_merging(self) -> List[str]:
+        """Construct merge query for MERGE mode."""
+        merge_table = f"{self.full_table_name}_merge"
+        quote = self.connector.quote_identifier
+        join = " AND ".join([f"a.{quote(x)} = b.{quote(x)}" for x in self.primary_keys])
+
+        drop_merge_table = self._make_drop_table_query(self.schema, merge_table)
+        queries = f"""
+            {drop_merge_table};
+
+            CREATE TABLE {quote(merge_table)} WITH (DISTRIBUTION = ROUND_ROBIN)
+            AS
+            SELECT a.*
+            FROM {self.quoted_full_table_name} AS a
+            LEFT JOIN {self.quoted_full_staging_table} AS b ON {join}
+            WHERE b.{quote(self.primary_keys[0])} IS NULL
+            UNION ALL
+            SELECT * FROM {self.quoted_full_staging_table};
+
+            TRUNCATE TABLE {self.quoted_full_table_name};
+            INSERT INTO {self.quoted_full_table_name} SELECT * FROM {quote(merge_table)};
+
+            {drop_merge_table};
+        """
+        return queries.split(";")

recurvedata/pigeon/loader/csv_to_mssql.py

@@ -0,0 +1,174 @@
+from recurvedata.pigeon import const
+from recurvedata.pigeon.connector.mssql import SQLServerConnector
+from recurvedata.pigeon.loader.base import BaseLoader, CSVToDBAPIMixin
+from recurvedata.pigeon.utils import ensure_query_list, ensure_str_list, fs
+
+allowed_modes = (const.LOAD_OVERWRITE, const.LOAD_MERGE, const.LOAD_APPEND)
+STATING_TABLE_NAME_PLACEHOLDER = "<TABLE>"
+
+
+class CSVToMsSQLLoader(BaseLoader, CSVToDBAPIMixin):
+    def __init__(
+        self,
+        database,
+        table,
+        filename,
+        connector: SQLServerConnector,
+        schema=None,
+        create_table_ddl=None,
+        staging_create_table_ddl=None,
+        mode=const.LOAD_OVERWRITE,
+        primary_keys=None,
+        skiprows=0,
+        columns=None,
+        using_insert=True,
+        insert_batch_size=500,
+        insert_concurrency=1,
+        delete_file=False,
+        pre_queries=None,
+        post_queries=None,
+        *args,
+        **kwargs,
+    ):
+        self.database = database
+        self.table = table
+
+        if "." in table:
+            self.schema, self.table = table.split(".")
+        else:
+            self.schema = schema or "dbo"
+            self.table = table
+
+        connector.database = self.database
+        self.connector = connector
+        self.filename = filename
+        self.create_table_ddl = create_table_ddl
+        # 考虑到完整的复制表结构（包括约束和索引）比较复杂，允许指定 staging 表的 DDL
+        # 表名用特殊符号 <TABLE> 占位
+        self.staging_create_table_ddl = staging_create_table_ddl
+        if self.staging_create_table_ddl and STATING_TABLE_NAME_PLACEHOLDER not in self.staging_create_table_ddl:
+            raise ValueError(f"use {STATING_TABLE_NAME_PLACEHOLDER} as table name placeholder")
+
+        if mode not in allowed_modes:
+            raise ValueError("mode should be one of ({})".format(allowed_modes))
+
+        self.mode = mode
+        self.primary_keys = ensure_str_list(primary_keys)
+        if self.mode == const.LOAD_MERGE and not self.primary_keys:
+            raise ValueError("primary_keys should not be empty in mode {}".format(const.LOAD_MERGE))
+
+        # self.columns = columns or self.csvfile.header
+        # self.skiprows = int(skiprows or self.csvfile.has_header)
+        self.columns = columns
+        self.skiprows = int(skiprows)
+        self.using_insert = using_insert
+        self.insert_batch_size = insert_batch_size
+        self.insert_concurrency = insert_concurrency
+        self.delete_file = delete_file
+
+        self.pre_queries = ensure_query_list(pre_queries) or []
+        self.post_queries = ensure_query_list(post_queries) or []
+
+        super().__init__()
+
+    @property
+    def staging_table(self):
+        return f"{self.schema}.{self.table}_staging"
+
+    @property
+    def full_table_name(self):
+        return f"{self.schema}.{self.table}"
+
+    @property
+    def quoted_staging_table(self):
+        return self.connector.quote_identifier(self.staging_table)
+
+    @property
+    def quoted_full_table_name(self):
+        return self.connector.quote_identifier(self.full_table_name)
+
+    def execute_impl(self):
+        if fs.is_file_empty(self.filename):
+            self.logger.error("file not exists or has no content. %s", self.filename)
+            fs.remove_files_safely(fs.schema_filename(self.filename))
+            return
+
+        self._prepare_target_table()
+        self._prepare_staging_table()
+        self._load_to_staging()
+        self._merge_into_target_table()
+
+        # do cleaning things
+        if self.delete_file:
+            self.logger.info("delete local file %s", self.filename)
+            fs.remove_files_safely(self.filename)
+            fs.remove_files_safely(fs.schema_filename(self.filename))
+
+    def _prepare_staging_table(self):
+        schema, table = self.staging_table.split(".")
+        if self.staging_create_table_ddl:
+            ddl: str = self.staging_create_table_ddl.replace(STATING_TABLE_NAME_PLACEHOLDER, self.quoted_staging_table)
+            ddl = ddl.rstrip(";")
+        else:
+            ddl = f"SELECT TOP 0 * INTO {self.quoted_staging_table} FROM {self.quoted_full_table_name}"
+
+        query = f"""
+            IF EXISTS (
+              SELECT * FROM sys.tables
+              WHERE schema_name(schema_id) = '{schema}' AND name = '{table}'
+            )
+            DROP TABLE {self.quoted_staging_table};
+
+            {ddl}
+        """
+        self.connector.execute(query)
+
+    def _load_to_staging(self):
+        self.logger.info("load %s into staging table %s", self.filename, self.staging_table)
+        self.connector.load_csv(
+            table=self.staging_table,
+            filename=self.filename,
+            schema=self.schema,
+            columns=self.columns,
+            skiprows=self.skiprows,
+            using_insert=self.using_insert,
+            null_values=("NULL", r"\N", ""),
+            batch_size=self.insert_batch_size,
+            concurrency=self.insert_concurrency,
+        )
+
+    def _merge_into_target_table(self):
+        target = self.quoted_full_table_name
+        staging = self.quoted_staging_table
+
+        queries = []
+        if self.mode == const.LOAD_OVERWRITE:
+            queries.append(f"TRUNCATE TABLE {target}")
+            append_sql = f"INSERT INTO {target} SELECT * FROM {staging}"
+            queries.append(append_sql)
+        elif self.mode == const.LOAD_MERGE:
+            joins = []
+            for field in self.primary_keys:
+                field = self.connector.quote_identifier(field)
+                join = f"{target}.{field} = {staging}.{field}"
+                joins.append(join)
+
+            join_conditions = " AND ".join(joins)
+            # Delete existing records that match primary keys
+            delete_sql = f"DELETE {target} FROM {target} INNER JOIN {staging} ON {join_conditions}"
+            queries.append(delete_sql)
+
+            # Insert all data from staging table to target table
+            insert_sql = f"INSERT INTO {target} SELECT * FROM {staging}"
+            queries.append(insert_sql)
+        else:
+            # APPEND mode
+            append_sql = f"INSERT INTO {target} SELECT * FROM {staging}"
+            queries.append(append_sql)
+
+        queries.append(f"DROP TABLE {staging}")
+
+        queries = self.pre_queries + queries + self.post_queries
+        self.logger.info("running SQL Server queries...")
+        self.connector.execute(queries, autocommit=False, commit_on_close=True)
+        self.logger.info("done.")

recurvedata/pigeon/loader/csv_to_mysql.py

@@ -0,0 +1,180 @@
+from recurvedata.pigeon import const
+from recurvedata.pigeon.connector import new_mysql_connector
+from recurvedata.pigeon.csv import CSV
+from recurvedata.pigeon.loader.base import BaseLoader, CSVToDBAPIMixin
+from recurvedata.pigeon.utils import ensure_query_list, ensure_str_list, fs
+from recurvedata.pigeon.utils.sql import bak_table_of, staging_table_of
+
+allowed_modes = (const.LOAD_OVERWRITE, const.LOAD_MERGE, const.LOAD_APPEND)
+
+
+class CSVToMySQLLoader(BaseLoader, CSVToDBAPIMixin):
+    def __init__(
+        self,
+        database,
+        table,
+        filename,
+        connector=None,
+        create_table_ddl=None,
+        mode=const.LOAD_OVERWRITE,
+        primary_keys=None,
+        skiprows=0,
+        columns=None,
+        using_insert=False,
+        insert_batch_size=1000,
+        insert_concurrency=1,
+        delete_file=False,
+        tidb_dml_batch_size=500,
+        pre_queries=None,
+        post_queries=None,
+        *args,
+        **kwargs,
+    ):
+        self.database = database
+        self.table = table
+
+        if isinstance(filename, CSV):
+            filename = filename.path
+        self.filename = filename
+        self.csvfile = CSV(self.filename)
+
+        if connector is None:
+            connector = new_mysql_connector(database=self.database)
+        else:
+            connector.database = self.database
+        self.connector = connector
+
+        self.create_table_ddl = create_table_ddl
+
+        if mode not in allowed_modes:
+            raise ValueError("mode should be one of ({})".format(allowed_modes))
+
+        self.mode = mode
+        self.primary_keys = ensure_str_list(primary_keys)
+        if self.mode == const.LOAD_MERGE and not self.primary_keys:
+            raise ValueError("primary_keys should not be empty in mode {}".format(const.LOAD_MERGE))
+
+        # self.columns = columns or self.csvfile.header
+        # self.skiprows = int(skiprows or self.csvfile.has_header)
+        self.columns = columns
+        self.skiprows = int(skiprows)
+        self.using_insert = using_insert
+        self.insert_batch_size = insert_batch_size
+        self.insert_concurrency = insert_concurrency
+        self.delete_file = delete_file
+
+        # https://pingcap.com/docs-cn/sql/tidb-specific/#tidb-dml-batch-size
+        self.tidb_dml_batch_size = tidb_dml_batch_size
+
+        self.pre_queries = ensure_query_list(pre_queries) or []
+        self.post_queries = ensure_query_list(post_queries) or []
+
+        super().__init__()
+
+    @property
+    def staging_table(self):
+        return staging_table_of(self.table)
+
+    def execute_impl(self):
+        if fs.is_file_empty(self.filename):
+            self.logger.error("file not exists or has no content. %s", self.filename)
+            fs.remove_files_safely(fs.schema_filename(self.filename))
+            return
+        self._prepare_target_table()
+        self._prepare_staging_table()
+        self._load_to_staging()
+        self._merge_into_target_table()
+
+        # do cleaning things
+        if self.delete_file:
+            self.logger.info("delete local file %s", self.filename)
+            fs.remove_files_safely(self.filename)
+            fs.remove_files_safely(fs.schema_filename(self.filename))
+
+    def _prepare_staging_table(self):
+        queries = """
+            DROP TABLE IF EXISTS {staging};
+            CREATE TABLE {staging} LIKE {table};
+        """.format(
+            staging=self.staging_table, table=self.table
+        )
+        self.connector.execute(queries, autocommit=True)
+
+    def _load_to_staging(self):
+        self.connector.load_csv(
+            table=self.staging_table,
+            filename=self.csvfile.path,
+            columns=self.columns,
+            lineterminator=self.csvfile.dialect.lineterminator,
+            skiprows=self.skiprows,
+            using_insert=self.using_insert,
+            null_values=("NULL", r"\N", ""),
+            batch_size=self.insert_batch_size,
+            concurrency=self.insert_concurrency,
+        )
+
+    def _merge_into_target_table(self):
+        queries = []
+        if self.connector.is_tidb():
+            queries.append("SET autocommit=1")
+            queries.append("SET @@session.tidb_batch_delete=ON")
+            queries.append("SET @@session.tidb_batch_insert=ON")
+            if self.tidb_dml_batch_size:
+                queries.append(f"SET @@session.tidb_dml_batch_size={self.tidb_dml_batch_size}")
+
+        if self.mode == const.LOAD_MERGE:
+            queries.extend(self._ingest_by_merging())
+        elif self.mode == const.LOAD_OVERWRITE:
+            bak_table = bak_table_of(self.table)
+            queries.append(f"DROP TABLE IF EXISTS {bak_table}")
+            queries.append(f"RENAME TABLE {self.table} TO {bak_table}")
+            queries.append(f"RENAME TABLE {self.staging_table} TO {self.table}")
+            queries.append(f"DROP TABLE IF EXISTS {bak_table}")
+        else:
+            queries.append(f"INSERT INTO {self.table} SELECT * FROM {self.staging_table}")
+            queries.append(f"DROP TABLE {self.staging_table}")
+
+        queries = self.pre_queries + queries + self.post_queries
+        self.logger.info("running MySQL queries within a transaction")
+        self.connector.execute(queries, autocommit=False, commit_on_close=True)
+
+    def _ingest_by_merging(self):
+        """Merge with deduplication based on specified primary_keys"""
+        # First, deduplicate staging table based on primary_keys using window function
+        pk_columns = ", ".join(self.primary_keys)
+
+        # Get all columns from staging table (excluding the rn column we'll add)
+        cols = self.connector.get_columns(self.staging_table)
+        cols_str = ", ".join(self.connector.quote_identifier(x) for x in cols)
+
+        # Create a temporary table with deduplicated data
+        tmp_table = f"{self.staging_table}_dedup"
+        dedup_sql = f"""
+            DROP TABLE IF EXISTS {tmp_table};
+            CREATE TABLE {tmp_table} LIKE {self.staging_table};
+            INSERT INTO {tmp_table}
+            SELECT {cols_str} FROM (
+                SELECT *, ROW_NUMBER() OVER(PARTITION BY {pk_columns} ORDER BY {pk_columns}) AS rn
+                FROM {self.staging_table}
+            ) t
+            WHERE rn = 1;
+        """
+
+        # Replace staging table with deduplicated data
+        replace_sql = f"""
+            DROP TABLE {self.staging_table};
+            RENAME TABLE {tmp_table} TO {self.staging_table};
+        """
+
+        # Delete records from target table that have the same primary keys as staging table
+        join_condition = " AND ".join([f"a.{pk} = b.{pk}" for pk in self.primary_keys])
+        delete_sql = f"""
+            DELETE a FROM {self.table} a
+            INNER JOIN {self.staging_table} b ON {join_condition}
+        """
+
+        # Insert deduplicated data into target table
+        insert_sql = f"INSERT INTO {self.table} SELECT * FROM {self.staging_table}"
+        drop_sql = f"DROP TABLE {self.staging_table}"
+
+        return [dedup_sql, replace_sql, delete_sql, insert_sql, drop_sql]