count_lines param

Author: Redevil10

README.md

@@ -96,6 +96,7 @@ CsvToPostgresOperator(
 | `has_header` | Include CSV header row | `True` |
 | `compression` | Compression format (`"gzip"` or `None`) | `None` |
 | `timeout` | Query timeout in minutes | `60` |
+| `count_lines` | Count and log the number of lines in the CSV after writing | `True` |
 
 ### CsvToPostgresOperator
 
@@ -112,6 +113,7 @@ CsvToPostgresOperator(
 | `quote_char` | CSV quote character | `'"'` |
 | `null_string` | String representing NULL | `""` |
 | `timeout` | Query timeout in minutes | `60` |
+| `count_lines` | Count and log the number of lines in the CSV before loading | `True` |
 
 ## Development
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "airflow-postgres-csv"
-version = "0.3.1"
+version = "0.3.2"
 description = "Airflow operators for PostgreSQL <-> CSV file transfers using COPY"
 readme = "README.md"
 license = "MIT"

src/airflow_postgres_csv/__init__.py

@@ -3,4 +3,4 @@
 from airflow_postgres_csv.operators import CsvToPostgresOperator, PostgresToCsvOperator
 
 __all__ = ["PostgresToCsvOperator", "CsvToPostgresOperator"]
-__version__ = "0.3.1"
+__version__ = "0.3.2"

src/airflow_postgres_csv/operators.py

@@ -34,6 +34,8 @@ class PostgresToCsvOperator(BaseOperator):
     :param compression: Compression format. Supports ``"gzip"``.
         Defaults to ``None`` (no compression).
     :param timeout: Query timeout in minutes. Defaults to ``60``.
+    :param count_lines: Count and log the number of lines in the CSV file after writing.
+        Disable for very large files to avoid the extra read. Defaults to ``True``.
     """
 
     template_fields: Sequence[str] = (
@@ -51,6 +53,7 @@ def __init__(
         has_header: bool = True,
         compression: str | None = None,
         timeout: int = 60,
+        count_lines: bool = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -61,6 +64,7 @@ def __init__(
         self.has_header = has_header
         self.compression = compression
         self.timeout = timeout
+        self.count_lines = count_lines
 
     def execute(self, context):
         sql = self.sql
@@ -89,15 +93,16 @@ def execute(self, context):
                     cursor.copy_expert(copy_command, csv_file)
                     rows = cursor.rowcount
 
-        open_func = self._get_open_func()
-        with open_func(self.csv_file_path, "rt", encoding="utf-8") as f:
-            line_count = sum(1 for _ in f)
+        if self.count_lines:
+            open_func = self._get_open_func()
+            with open_func(self.csv_file_path, "rt", encoding="utf-8") as f:
+                line_count = sum(1 for _ in f)
 
         self.log.info(
-            "CSV saved: %s (%s rows, %s lines, %s)",
+            "CSV saved: %s (%s rows%s, %s)",
             self.csv_file_path,
             rows if rows >= 0 else "unknown",
-            line_count,
+            f", {line_count} lines" if self.count_lines else "",
             "with header" if self.has_header else "no header",
         )
         return self.csv_file_path
@@ -126,6 +131,8 @@ class CsvToPostgresOperator(BaseOperator):
     :param compression: Compression format. Supports ``"gzip"``.
         Defaults to ``None`` (no compression).
     :param timeout: Query timeout in minutes. Defaults to ``60``.
+    :param count_lines: Count and log the number of lines in the CSV file before loading.
+        Disable for very large files to avoid the extra read. Defaults to ``True``.
     """
 
     template_fields: Sequence[str] = ("csv_file_path", "table_name")
@@ -143,6 +150,7 @@ def __init__(
         truncate: bool = False,
         compression: str | None = None,
         timeout: int = 60,
+        count_lines: bool = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -157,21 +165,23 @@ def __init__(
         self.truncate = truncate
         self.compression = compression
         self.timeout = timeout
+        self.count_lines = count_lines
 
     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)
 
-        open_func = self._get_open_func()
-        with open_func(self.csv_file_path, "rt", encoding="utf-8") as f:
-            line_count = sum(1 for _ in f)
+        if self.count_lines:
+            open_func = self._get_open_func()
+            with open_func(self.csv_file_path, "rt", encoding="utf-8") as f:
+                line_count = sum(1 for _ in f)
 
         self.log.info(
-            "Loading %s (%s lines, %s) into %s",
+            "Loading %s (%s%s) into %s",
             self.csv_file_path,
-            line_count,
+            f"{line_count} lines, " if self.count_lines else "",
             "with header" if self.has_header else "no header",
             self.table_name,
         )