SCUDO 0.1.0__py3-none-any.whl

cudo/doris/client.py

@@ -0,0 +1,310 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+'''
+# Doris Flight SQL Client (ADBC)
+
+Python client for interacting with Apache Doris using Arrow Flight SQL via ADBC.
+
+This class provides:
+- Simple connection handling
+- Query execution
+- Optimized DataFrame fetching (`fetch_df`)
+- Optional Arrow-based and partition-based methods
+- Timing and logging utilities
+
+## Requirements
+
+```bash
+pip install adbc_driver_manager adbc_driver_flightsql pandas
+```
+
+## Example
+    client = DorisFlightSQLClient(
+        uri="grpc://127.0.0.1:9090",
+        username="root",
+        password=""
+    )
+
+    df = client.query("SELECT * FROM my_table LIMIT 1000;")
+    print(df.head())
+
+    client.close()
+'''
+
+import adbc_driver_manager
+import adbc_driver_flightsql.dbapi as flight_sql
+import pandas as pd
+from datetime import datetime
+from typing import Optional
+
+
+
+class DorisFlightSQLClient:
+    '''
+    Apache Doris Flight SQL client using ADBC (DB-API 2.0).
+    This client wraps the ADBC Flight SQL driver and provides optimized
+    methods for querying data into pandas DataFrames.
+
+    Parameters
+    ----------
+    uri : str
+        Flight SQL endpoint (e.g., "grpc://127.0.0.1:9090").
+
+    username : str
+        Database username.
+
+    password : str
+        Database password.
+
+    Attributes
+    ----------
+    conn : flight_sql.Connection
+        Active database connection.
+
+    cursor : flight_sql.Cursor
+        Cursor used for executing queries.
+
+    Notes
+    -----
+    - Uses `fetch_df()` by default for best performance.
+    - Compatible with Apache Doris Arrow Flight SQL service.
+    '''
+
+    def __init__(self, uri: str, username: str, password: str) -> None:
+        self.uri = uri
+        self.db_kwargs = {
+            adbc_driver_manager.DatabaseOptions.USERNAME.value: username,
+            adbc_driver_manager.DatabaseOptions.PASSWORD.value: password,
+        }
+
+        self.conn: Optional[flight_sql.Connection] = None
+        self.cursor: Optional[flight_sql.Cursor] = None
+
+        self._connect()
+
+
+    # ------------------------------------------------------------------
+    # Connection Handling
+    # ------------------------------------------------------------------
+
+    def _connect(self) -> None:        
+        '''
+        Establish a connection and create a cursor.
+        '''
+        try:
+            self.conn = flight_sql.connect(uri=self.uri, db_kwargs=self.db_kwargs)
+            self.cursor = self.conn.cursor()
+        except Exception as e:
+            raise RuntimeError("Connection failed") from e
+
+    def close(self) -> None:
+        '''
+        Close cursor and connection.
+        '''
+        try:
+            if self.cursor:
+                self.cursor.close()
+        finally:
+            if self.conn:
+                self.conn.close()
+
+
+    # ------------------------------------------------------------------
+    # Core Query Method
+    # ------------------------------------------------------------------
+
+    def query(self, sql: str, params: dict | None = None) -> pd.DataFrame:
+        '''
+        Execute a SQL query and return a pandas DataFrame.
+
+        This method uses `fetch_df()` (ADBC optimized path),
+        which is faster than converting from Arrow manually.
+
+        Parameters
+        ----------
+        sql : str
+            SQL query string.
+
+        params : dict, optional
+            Parameters to safely inject into SQL.
+
+        Returns
+        -------
+        pandas.DataFrame
+            Query result as a DataFrame.
+
+        Example
+        -------
+        >>> df = client.query("SELECT * FROM table LIMIT 100;")
+        
+        '''
+        if not self.cursor:
+            raise RuntimeError("Cursor is not initialized.")
+
+        sql = self._render_sql(sql, params=params)
+
+        start_time = datetime.now()
+
+        try:
+            self.cursor.execute(sql)
+            df = self.cursor.fetch_df()
+        except Exception as e:
+            raise RuntimeError("Query execution failed") from e
+
+        duration = datetime.now() - start_time
+
+        print("\n##################")
+        print(f"Query executed in: {duration}")
+        print(df.info(memory_usage="deep"))
+
+        return df
+
+
+    # ------------------------------------------------------------------
+    # Alternative Methods
+    # ------------------------------------------------------------------
+
+    def query_arrow(self, sql: str, params: dict | None = None) -> pd.DataFrame:
+        '''
+        Execute query using Arrow fetch, then convert to pandas.
+
+        Slower than `query()` but useful for debugging or Arrow workflows.
+
+        Parameters
+        ----------
+        sql : str
+            SQL query.
+
+        params : dict, optional
+            Parameters to safely inject into SQL.
+
+        Returns
+        -------
+        pandas.DataFrame
+        '''
+        sql = self._render_sql(sql, params=params)
+
+        start_time = datetime.now()
+
+        try:
+            self.cursor.execute(sql)
+            arrow_data = self.cursor.fetchallarrow()
+            df = arrow_data.to_pandas()
+        except Exception as e:
+            raise RuntimeError("Arrow query execution failed") from e
+
+        duration = datetime.now() - start_time
+
+        print("\n##################")
+        print(f"Arrow fetch in: {duration}, bytes: {arrow_data.nbytes}")
+        print(df.info(memory_usage="deep"))
+
+        return df
+
+
+    def query_partition(self, sql: str, params: dict | None = None) -> pd.DataFrame:
+        '''
+        Execute query using partitioned reads.
+
+        Useful for parallel processing of large datasets.
+
+        Parameters
+        ----------
+        sql : str
+            SQL query.
+            
+        params : dict, optional
+            Parameters to safely inject into SQL.
+
+        Returns
+        -------
+        pandas.DataFrame
+        '''
+        sql = self._render_sql(sql, params=params)
+
+        start_time = datetime.now()
+
+        try:
+            partitions, schema = self.cursor.adbc_execute_partitions(sql)
+            # Example: read first partition (extendable to parallel processing)
+            self.cursor.adbc_read_partition(partitions[0])
+            arrow_data = self.cursor.fetchallarrow()
+            df = arrow_data.to_pandas()
+        except Exception as e:
+            raise RuntimeError("Partition query execution failed") from e
+
+        duration = datetime.now() - start_time
+
+        print("\n##################")
+        print(f"Partitions: {len(partitions)}, time: {duration}")
+        print(df.info(memory_usage="deep"))
+
+        return df
+
+
+    # ------------------------------------------------------------------
+    # Utility Methods
+    # ------------------------------------------------------------------
+
+    def _render_sql(self, sql: str, params: dict | None = None) -> str:
+        '''
+        Render SQL using simple parameter substitution.
+
+        Parameters
+        ----------
+        sql : str
+            SQL query with `{param}` placeholders.
+            
+        params : dict, optional
+            Dictionary of parameters.
+
+        Returns
+        -------
+        str
+            Rendered SQL string.
+        '''
+        if params:
+            safe_params = {}
+            for k, v in params.items():
+                if isinstance(v, (int, float)):
+                    safe_params[k] = v
+                elif isinstance(v, str):
+                    safe_params[k] = "'" + v.replace("'", "''") + "'"
+                else:
+                    raise ValueError(f"Unsafe parameter type for {k}: {type(v)}")
+        try:
+            return sql.format(**safe_params)
+        except KeyError as e:
+            raise ValueError(f"Missing SQL parameter: {e}") from e
+
+
+    def execute(self, sql: str) -> None:
+        '''
+        Execute a SQL statement without returning results.
+
+        Useful for DDL or INSERT operations.
+
+        Parameters
+        ----------
+        sql : str
+        '''
+        if not self.cursor:
+            raise RuntimeError("Cursor is not initialized.")
+
+        print(f"\n### Executing SQL ###\n{sql}")
+        self.cursor.execute(sql)
+
+
+    def __enter__(self):
+        ''' Context manager entry.'''
+        return self
+
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        ''' Context manager exit (auto-close).'''
+        if exc_type:
+            print(f"[DorisFlightSQLClient] Exception: {exc_val}")
+        self.close()
+        return False  # always propagate errors
+

cudo/doris/pymysql.py

@@ -0,0 +1,380 @@
+'''
+Doris Client Module
+==============================
+
+
+This module provides a lightweight client for interacting with **Apache Doris**
+using a MySQL-compatible interface.
+
+It encapsulates common operations required for data loading workflows:
+- Establishing a connection to Doris
+- Submitting LOAD jobs
+- Monitoring job execution status
+- Waiting for job completion
+
+The implementation is designed to be simple, reusable, and easily integrable
+into ETL/ELT pipelines.
+
+Main Features
+-------------
+
+- Connection management using `pymysql`
+- Load job submission via SQL
+- Polling-based job monitoring
+- Optional logging integration
+
+Example Usage
+-------------
+
+```python
+client = DorisClient(config, logger)
+
+client.submit_load("my_db", load_sql)
+client.wait_for_load("my_label")
+```
+
+'''
+
+import time
+import pymysql
+
+
+class DorisPyMySQLClient:
+    '''
+    Client for interacting with Apache Doris.
+
+    This class provides methods to execute and monitor data load operations
+    in Apache Doris using its MySQL-compatible interface.
+
+    Args:
+        config (dict):
+            Configuration dictionary containing connection parameters:
+            - host (str): Doris FE host
+            - port (int): Doris FE MySQL port
+            - user (str): Username
+            - password (str): Password
+            - autocommit (bool): Autocommit flag
+
+        logger (logging.Logger, optional):
+            Logger instance for logging execution details.
+
+    Attributes:
+        conn (pymysql.Connection):
+            Active database connection.
+    '''
+    def __init__(self, config, logger=None):
+        self.logger = logger
+        self.conn = self._connect(config)
+
+    def _connect(self, config):
+        '''
+        Establish a connection to Apache Doris.
+
+        Args:
+            config (dict):
+                Dictionary containing connection parameters.
+
+        Returns:
+            pymysql.Connection:
+                Active connection to Doris.
+
+        Raises:
+            RuntimeError:
+                If the connection cannot be established.
+        '''
+        try:
+            conn = pymysql.connect(
+                host=config["host"],
+                port=config["port"],
+                user=config["user"],
+                password=config["password"],
+                autocommit=config["autocommit"]
+            )
+            return conn
+
+        except Exception as e:
+            if self.logger:
+                self.logger.error("Connection failed: %s", str(e))
+            raise RuntimeError("Module failed")
+
+    def close(self):
+        '''
+        Close Doris connection.
+        '''
+        if self.conn:
+            self.conn.close()
+
+    def submit_load(self, db_name, sql) -> None:
+        '''
+        Submit a LOAD job to Apache Doris.
+
+        This method switches to the target database and executes the provided
+        LOAD SQL statement.
+
+        Args:
+            db_name (str):
+                Target database name.
+
+            sql (str):
+                Full Doris LOAD SQL statement.
+
+        Returns:
+            None
+
+        Raises:
+            RuntimeError:
+                If the SQL execution fails.
+        '''
+        try:
+            with self.conn.cursor() as cur:
+                cur.execute(f"USE {db_name}")
+                cur.execute(sql)
+
+            if self.logger:
+                self.logger.info("Load submitted successfully")
+
+        except Exception as e:
+            if self.logger:
+                self.logger.error("Submit load failed: %s", str(e))
+            raise RuntimeError("Module failed")
+
+    def check_load_status(self, label) -> str:
+        '''
+        Check the status of a Doris LOAD job.
+
+        Args:
+            label (str):
+                Unique label identifying the load job.
+
+        Returns:
+            str | None:
+                - Job state (e.g., ``FINISHED``, ``CANCELLED``)
+                - ``None`` if the job is not yet visible
+
+        Raises:
+            RuntimeError:
+                If the query execution fails.
+        '''
+        try:
+            with self.conn.cursor() as cur:
+                cur.execute(f"SHOW LOAD WHERE Label = '{label}'")
+                result = cur.fetchone()
+
+            if not result:
+                return None
+
+            return result[2]  # state column
+
+        except Exception as e:
+            if self.logger:
+                self.logger.error("Check load status failed: %s", str(e))
+            raise RuntimeError("Module failed")
+
+    def wait_for_load(self, label) -> None:
+        '''
+        Wait until a Doris LOAD job completes.
+
+        This method continuously polls the job status until:
+        - The job finishes successfully (``FINISHED``)
+        - The job fails (``CANCELLED``)
+
+        Args:
+            label (str):
+                Load job label.
+
+        Returns:
+            None
+
+        Raises:
+            RuntimeError:
+                If the load job fails or is cancelled.
+        '''
+        try:
+            while True:
+                state = self.check_load_status(label)
+
+                if state is None:
+                    if self.logger:
+                        self.logger.info("Load job not visible yet...")
+
+                elif state == "FINISHED":
+                    if self.logger:
+                        self.logger.info("Load job finished successfully")
+                    return
+
+                elif state == "CANCELLED":
+                    if self.logger:
+                        self.logger.error("Doris load job failed")
+                    raise RuntimeError("Doris load job failed")
+
+                else:
+                    if self.logger:
+                        self.logger.info(f"Current state: {state}")
+
+                time.sleep(5)
+
+        except Exception as e:
+            if self.logger:
+                self.logger.error("Wait for load failed: %s", str(e))
+            raise RuntimeError("Module failed")
+
+
+
+# import time
+# import pymysql
+
+
+# def connect_database(config, logger=None):
+#     '''
+#     Create a database connection to Doris
+#     '''
+
+#     try:
+#         conn = pymysql.connect(
+#             host=config["host"],
+#             port=config["port"],
+#             user=config["user"],
+#             password=config["password"],
+#             autocommit=config["autocommit"]
+#         )
+#         return conn
+
+#     except Exception as e:
+#         if logger: logger.error("module failed: %s", str(e))
+#         raise RuntimeError("Module failed")
+
+# # to call:
+# # wait_for_parquet(client,args.bucket, args.parquet_prefix )
+# # 
+# # def wait_for_parquet(client, bucket, prefix, timeout=600):
+
+# #     start = time.time()
+
+# #     while True:
+
+# #         objects = list(client.list_objects(bucket, prefix=prefix, recursive=True))
+# #         parquet_files = [o.object_name for o in objects if o.object_name.endswith(".parquet")]
+
+# #         if parquet_files:
+# #             print(f"Found {len(parquet_files)} parquet files")
+# #             return parquet_files
+
+# #         if time.time() - start > timeout:
+# #             raise RuntimeError("Timeout waiting for parquet files")
+
+# #         print("No parquet files yet, sleeping 10s...")
+# #         time.sleep(10)
+
+
+# def submit_load(conn, db_name, sql, logger=None):
+#     '''
+#     Submit a LOAD job to Apache Doris.
+
+#     Args:
+#         conn (pymysql.Connection):
+#             Active MySQL-compatible connection to Doris.
+
+#         db_name (str):
+#             Target database name.
+
+#         sql (str):
+#             Complete Doris LOAD SQL statement.
+
+#     Returns:
+#         None
+
+#     Raises:
+#         pymysql.MySQLError:
+#             If the SQL execution fails.
+#     '''
+
+#     try:
+#         with conn.cursor() as cur:
+#             cur.execute(f"USE {db_name}")
+#             cur.execute(sql)
+
+#     except Exception as e:
+#         if logger: logger.error("module failed: %s", str(e))
+#         raise RuntimeError("Module failed")
+
+
+# def check_load_status(conn, label, logger=None):
+#     '''
+#     Retrieve the current state of a Doris LOAD job.
+
+#     Args:
+#         conn (pymysql.Connection):
+#             Active MySQL-compatible connection to Doris.
+
+#         label (str):
+#             Load job label used to identify the job.
+
+#     Returns:
+#         str | None:
+#             Current state of the job (e.g. ``FINISHED`` or ``CANCELLED``),
+#             or ``None`` if the job is not yet visible.
+#     '''
+
+#     try:
+#         with conn.cursor() as cur:
+#             cur.execute(f"SHOW LOAD WHERE Label = '{label}'")
+#             result = cur.fetchone()
+
+#         if not result:
+#             return None
+
+#         return result[2]  # state column
+
+#     except Exception as e:
+#         if logger: logger.error("module failed: %s", str(e))
+#         raise RuntimeError("Module failed")
+
+
+
+# def wait_for_load(conn, label, logger=None):
+#     '''
+#     Wait until a Doris LOAD job finishes.
+
+#     This function repeatedly polls the job status until the load
+#     finishes successfully or fails.
+
+#     Args:
+#         conn (pymysql.Connection):
+#             Active MySQL-compatible connection to Doris.
+
+#         label (str):
+#             Label identifying the load job.
+
+#         logger (logging.Logger, optional):
+#             Logger used to report progress.
+
+#     Returns:
+#         None
+
+#     Raises:
+#         RuntimeError:
+#             If the load job is cancelled or fails.
+#     '''
+
+#     try:
+
+#         while True:
+
+#             state = check_load_status(conn, label)
+
+#             if state is None:
+#                 if logger: logger.info("load job not visible yet...")
+#             elif state == "FINISHED":
+#                 if logger: logger.info("load job finished successfully")
+#                 return
+#             elif state == "CANCELLED":
+#                 if logger: logger.error("Doris load job failed")
+#                 raise RuntimeError("Doris load job failed")
+#             else:
+#                 if logger: logger.info(f"Current state: {state}")
+
+#             time.sleep(5)
+
+#     except Exception as e:
+#         if logger: logger.error("module failed: %s", str(e))
+#         raise RuntimeError("Module failed")