airflow-postgres-csv 0.1.0__tar.gz

airflow_postgres_csv-0.1.0/.github/workflows/lint.yml

@@ -0,0 +1,19 @@
+name: Lint
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+      - name: Lint
+        run: uvx ruff check src/ tests/
+      - name: Format check
+        run: uvx ruff format --check src/ tests/

airflow_postgres_csv-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,44 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install build dependencies
+        run: pip install build
+
+      - name: Build package
+        run: python -m build
+
+      - name: Upload build artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - name: Download build artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

airflow_postgres_csv-0.1.0/.github/workflows/test.yml

@@ -0,0 +1,30 @@
+name: Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        run: uv python install ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          uv venv .venv --python ${{ matrix.python-version }}
+          uv pip install -e ".[dev]" --python .venv/bin/python
+
+      - name: Test
+        run: .venv/bin/pytest -v

airflow_postgres_csv-0.1.0/.gitignore

@@ -0,0 +1,17 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+dist/
+build/
+.eggs/
+.venv/
+.pytest_cache/
+.ruff_cache/
+.python-version
+*.csv
+
+# pycharm
+/.idea/
+
+# macOS
+**/.DS_Store

airflow_postgres_csv-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,7 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.9.6
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format

airflow_postgres_csv-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

airflow_postgres_csv-0.1.0/PKG-INFO

@@ -0,0 +1,95 @@
+Metadata-Version: 2.4
+Name: airflow-postgres-csv
+Version: 0.1.0
+Summary: Airflow operators for PostgreSQL <-> CSV file transfers using COPY
+Project-URL: Repository, https://github.com/Redevil10/airflow-postgres-csv
+Author: Qing Wan
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Framework :: Apache Airflow
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.10
+Requires-Dist: apache-airflow-providers-postgres>=6.0.0
+Requires-Dist: apache-airflow>=3.0.0
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# airflow-postgres-csv
+
+Airflow 3 operators for bulk PostgreSQL ↔ CSV transfers using `COPY`.
+
+## Operators
+
+- **`PostgresToCsvOperator`** — Run a SQL query and export results to a CSV file
+- **`CsvToPostgresOperator`** — Load a CSV file into a PostgreSQL table
+
+Both use PostgreSQL's `COPY` command for maximum throughput.
+
+## Installation
+
+```bash
+pip install airflow-postgres-csv
+```
+
+## Usage
+
+```python
+from airflow_postgres_csv import PostgresToCsvOperator, CsvToPostgresOperator
+
+# Export query results to CSV
+export_task = PostgresToCsvOperator(
+    task_id="export_users",
+    conn_id="my_postgres",
+    sql_query="SELECT * FROM users WHERE active = %(active)s",
+    query_params={"active": True},
+    csv_file_path="/tmp/users.csv",
+)
+
+# Load CSV into a table
+import_task = CsvToPostgresOperator(
+    task_id="import_users",
+    conn_id="my_postgres",
+    table_name="staging.users",
+    csv_file_path="/tmp/users.csv",
+)
+```
+
+### Key parameters
+
+**PostgresToCsvOperator:**
+| Parameter | Description | Default |
+|---|---|---|
+| `conn_id` | Airflow Postgres connection ID | required |
+| `csv_file_path` | Output file path (template-able) | required |
+| `sql_query` | SQL to execute | `None` |
+| `sql_file_path` | Path to `.sql` file | `None` |
+| `query_params` | Dict passed to `cursor.mogrify` | `{}` |
+| `has_header` | Include CSV header | `True` |
+| `timeout` | Query timeout (minutes) | `60` |
+
+**CsvToPostgresOperator:**
+| Parameter | Description | Default |
+|---|---|---|
+| `conn_id` | Airflow Postgres connection ID | required |
+| `table_name` | Target table (template-able, supports `schema.table`) | required |
+| `csv_file_path` | Input file path (template-able) | required |
+| `columns` | Explicit column list | `None` |
+| `has_header` | CSV has header row | `True` |
+| `delimiter` | CSV delimiter | `","` |
+| `quote_char` | CSV quote character | `'"'` |
+| `null_string` | String representing NULL | `""` |
+| `timeout` | Query timeout (minutes) | `60` |
+
+## Requirements
+
+- Apache Airflow >= 3.0.0
+- apache-airflow-providers-postgres >= 6.0.0
+
+## License
+
+MIT

airflow_postgres_csv-0.1.0/README.md

@@ -0,0 +1,74 @@
+# airflow-postgres-csv
+
+Airflow 3 operators for bulk PostgreSQL ↔ CSV transfers using `COPY`.
+
+## Operators
+
+- **`PostgresToCsvOperator`** — Run a SQL query and export results to a CSV file
+- **`CsvToPostgresOperator`** — Load a CSV file into a PostgreSQL table
+
+Both use PostgreSQL's `COPY` command for maximum throughput.
+
+## Installation
+
+```bash
+pip install airflow-postgres-csv
+```
+
+## Usage
+
+```python
+from airflow_postgres_csv import PostgresToCsvOperator, CsvToPostgresOperator
+
+# Export query results to CSV
+export_task = PostgresToCsvOperator(
+    task_id="export_users",
+    conn_id="my_postgres",
+    sql_query="SELECT * FROM users WHERE active = %(active)s",
+    query_params={"active": True},
+    csv_file_path="/tmp/users.csv",
+)
+
+# Load CSV into a table
+import_task = CsvToPostgresOperator(
+    task_id="import_users",
+    conn_id="my_postgres",
+    table_name="staging.users",
+    csv_file_path="/tmp/users.csv",
+)
+```
+
+### Key parameters
+
+**PostgresToCsvOperator:**
+| Parameter | Description | Default |
+|---|---|---|
+| `conn_id` | Airflow Postgres connection ID | required |
+| `csv_file_path` | Output file path (template-able) | required |
+| `sql_query` | SQL to execute | `None` |
+| `sql_file_path` | Path to `.sql` file | `None` |
+| `query_params` | Dict passed to `cursor.mogrify` | `{}` |
+| `has_header` | Include CSV header | `True` |
+| `timeout` | Query timeout (minutes) | `60` |
+
+**CsvToPostgresOperator:**
+| Parameter | Description | Default |
+|---|---|---|
+| `conn_id` | Airflow Postgres connection ID | required |
+| `table_name` | Target table (template-able, supports `schema.table`) | required |
+| `csv_file_path` | Input file path (template-able) | required |
+| `columns` | Explicit column list | `None` |
+| `has_header` | CSV has header row | `True` |
+| `delimiter` | CSV delimiter | `","` |
+| `quote_char` | CSV quote character | `'"'` |
+| `null_string` | String representing NULL | `""` |
+| `timeout` | Query timeout (minutes) | `60` |
+
+## Requirements
+
+- Apache Airflow >= 3.0.0
+- apache-airflow-providers-postgres >= 6.0.0
+
+## License
+
+MIT

airflow_postgres_csv-0.1.0/pyproject.toml

@@ -0,0 +1,46 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "airflow-postgres-csv"
+version = "0.1.0"
+description = "Airflow operators for PostgreSQL <-> CSV file transfers using COPY"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.10"
+authors = [{ name = "Qing Wan" }]
+classifiers = [
+    "Framework :: Apache Airflow",
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+]
+
+dependencies = [
+    "apache-airflow>=3.0.0",
+    "apache-airflow-providers-postgres>=6.0.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0",
+    "pytest-cov",
+    "ruff",
+    "pre-commit",
+]
+
+[project.urls]
+Repository = "https://github.com/Redevil10/airflow-postgres-csv"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/airflow_postgres_csv"]
+
+[tool.ruff]
+target-version = "py310"
+line-length = 99
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

airflow_postgres_csv-0.1.0/src/airflow_postgres_csv/__init__.py

@@ -0,0 +1,6 @@
+"""Airflow operators for PostgreSQL <-> CSV file transfers using COPY."""
+
+from airflow_postgres_csv.operators import CsvToPostgresOperator, PostgresToCsvOperator
+
+__all__ = ["PostgresToCsvOperator", "CsvToPostgresOperator"]
+__version__ = "0.1.0"

airflow_postgres_csv-0.1.0/src/airflow_postgres_csv/operators.py

@@ -0,0 +1,183 @@
+"""Custom Airflow operators for PostgreSQL <-> CSV file transfers."""
+
+import os
+from collections.abc import Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.providers.postgres.hooks.postgres import PostgresHook
+from airflow.sdk.bases.operator import BaseOperator
+
+
+class PostgresToCsvOperator(BaseOperator):
+    """
+    Execute a SQL query on PostgreSQL and save the result as a CSV file.
+
+    Uses ``COPY (...) TO STDOUT WITH CSV`` for high-performance bulk export.
+
+    :param conn_id: Airflow connection ID for the PostgreSQL database.
+    :param csv_file_path: Local file path where the CSV will be saved.
+    :param sql_query: SQL query to execute. Overrides *sql_file_path* if both given.
+    :param sql_file_path: Path to a ``.sql`` file containing the query.
+    :param query_params: Parameters passed to the SQL query via ``cursor.mogrify``.
+    :param has_header: Include a CSV header row. Defaults to ``True``.
+    :param timeout: Query timeout in minutes. Defaults to ``60``.
+    """
+
+    template_fields: Sequence[str] = (
+        "sql_query",
+        "sql_file_path",
+        "csv_file_path",
+    )
+
+    def __init__(
+        self,
+        conn_id: str,
+        csv_file_path: str,
+        sql_query: str | None = None,
+        sql_file_path: str | None = None,
+        query_params: dict | None = None,
+        has_header: bool = True,
+        timeout: int = 60,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.conn_id = conn_id
+        self.csv_file_path = csv_file_path
+        self.sql_query = sql_query
+        self.sql_file_path = sql_file_path
+        self.query_params = query_params or {}
+        self.has_header = has_header
+        self.timeout = timeout
+
+    def execute(self, context):
+        if not self.sql_query and not self.sql_file_path:
+            raise AirflowException("Either sql_query or sql_file_path must be provided")
+
+        if not self.sql_query:
+            with open(self.sql_file_path, encoding="utf-8") as f:
+                self.sql_query = f.read()
+
+        pg_hook = PostgresHook(postgres_conn_id=self.conn_id)
+        cleaned_sql = self.sql_query.strip().rstrip(";")
+
+        self.log.info("Running query and saving to CSV: %s", self.csv_file_path)
+
+        with pg_hook.get_conn() as conn:
+            with conn.cursor() as cursor:
+                cursor.execute("SET statement_timeout = %s;", (self.timeout * 60 * 1000,))
+                formatted_sql = cursor.mogrify(cleaned_sql, self.query_params).decode("utf-8")
+
+                header_clause = " HEADER" if self.has_header else ""
+                copy_command = f"COPY ({formatted_sql}) TO STDOUT WITH CSV{header_clause}"
+
+                rows = 0
+                with open(self.csv_file_path, "w", encoding="utf-8") as csv_file:
+                    cursor.copy_expert(copy_command, csv_file)
+                    rows = cursor.rowcount
+
+        self.log.info(
+            "CSV saved: %s (%s rows, %s)",
+            self.csv_file_path,
+            rows if rows >= 0 else "unknown",
+            "with header" if self.has_header else "no header",
+        )
+        return self.csv_file_path
+
+
+class CsvToPostgresOperator(BaseOperator):
+    """
+    Load a CSV file into a PostgreSQL table.
+
+    Uses ``COPY ... FROM STDIN WITH CSV`` for high-performance bulk import.
+
+    :param conn_id: Airflow connection ID for the PostgreSQL database.
+    :param table_name: Target table (may include schema, e.g. ``"myschema.mytable"``).
+    :param csv_file_path: Local file path of the CSV to load.
+    :param delimiter: CSV delimiter. Defaults to ``','``.
+    :param quote_char: CSV quote character. Defaults to ``'"'``.
+    :param null_string: String representing NULL values. Defaults to ``''``.
+    :param has_header: Whether the CSV has a header row. Defaults to ``True``.
+    :param columns: Explicit column list. If provided, maps CSV columns to these
+        table columns and skips the file header (if present).
+    :param timeout: Query timeout in minutes. Defaults to ``60``.
+    """
+
+    template_fields: Sequence[str] = ("csv_file_path", "table_name")
+
+    def __init__(
+        self,
+        conn_id: str,
+        table_name: str,
+        csv_file_path: str,
+        delimiter: str = ",",
+        quote_char: str = '"',
+        null_string: str = "",
+        has_header: bool = True,
+        columns: list[str] | None = None,
+        timeout: int = 60,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.conn_id = conn_id
+        self.table_name = table_name
+        self.csv_file_path = csv_file_path
+        self.delimiter = delimiter
+        self.quote_char = quote_char
+        self.null_string = null_string
+        self.has_header = has_header
+        self.columns = columns
+        self.timeout = timeout
+
+    def execute(self, context):
+        if not os.path.exists(self.csv_file_path):
+            raise AirflowException(f"CSV file not found: {self.csv_file_path}")
+
+        pg_hook = PostgresHook(postgres_conn_id=self.conn_id)
+
+        self.log.info("Loading %s into %s", self.csv_file_path, self.table_name)
+
+        column_clause = self._build_column_clause()
+        header_clause = "HEADER" if self.has_header and not self.columns else ""
+
+        copy_command = (
+            f"COPY {self._quote_table_name()} {column_clause} "
+            f"FROM STDIN WITH CSV "
+            f"DELIMITER '{self.delimiter}' "
+            f"QUOTE '{self.quote_char}' "
+            f"NULL '{self.null_string}' "
+            f"{header_clause}"
+        )
+
+        with pg_hook.get_conn() as conn:
+            with conn.cursor() as cursor:
+                cursor.execute("SET statement_timeout = %s;", (self.timeout * 60 * 1000,))
+                with open(self.csv_file_path, encoding="utf-8") as csv_file:
+                    if self.columns and self.has_header:
+                        next(csv_file)
+                    cursor.copy_expert(copy_command, csv_file)
+                rows = cursor.rowcount
+            conn.commit()
+
+        self.log.info(
+            "Loaded %s rows from %s into %s",
+            rows if rows >= 0 else "unknown",
+            self.csv_file_path,
+            self.table_name,
+        )
+        return rows
+
+    @staticmethod
+    def _quote_identifier(name: str) -> str:
+        """Quote a SQL identifier, escaping any double quotes."""
+        escaped = name.replace('"', '""')
+        return f'"{escaped}"'
+
+    def _quote_table_name(self) -> str:
+        parts = self.table_name.split(".")
+        return ".".join(self._quote_identifier(p) for p in parts)
+
+    def _build_column_clause(self) -> str:
+        if not self.columns:
+            return ""
+        cols = ", ".join(self._quote_identifier(c) for c in self.columns)
+        return f"({cols})"