[cdd_gae/parquet_to_table.py] Using binary COPY FROM via pgcopy (WiP my contributions to that package) ; [requirements.txt] Add pgcopy ; [cdd_gae/tests/test_parquet_to_table.py] Improve whitespace testing ;

Author: SamuelMarks

cdd_gae/parquet_to_table.py

@@ -5,14 +5,14 @@
 from collections import deque
 from datetime import datetime
 from functools import partial
-from io import StringIO
 from json import dumps, loads
 from operator import methodcaller
 from os import environ, path
 
 import numpy as np
 import psycopg2.sql
-from cdd.shared.pure_utils import identity, pp
+from cdd.shared.pure_utils import identity
+from pgcopy import CopyManager
 from pyarrow.parquet import ParquetFile
 from sqlalchemy import create_engine
 
@@ -70,6 +70,10 @@ def parse_col(col):
         raise NotImplementedError(type(col))
 
 
+# {"\\"first item\\"","\\"second item\\"",
+#  "\\"third item with \\\\\\"quotes\\\\\\" inside\\"","\\"fourth with a \\\\\\\\ backslash\\""}
+
+
 def maybe_quote_and_escape(s, ch='"'):
     if isinstance(s, str) and s.startswith("{") and s.endswith("}"):
         return "{ch}{}{ch}".format(s.replace('"', '\\"'), ch=ch)
@@ -88,44 +92,17 @@ def psql_insert_copy(table, conn, keys, data_iter):
         Column names
     data_iter : Iterable that iterates the values to be inserted
     """
-    # gets a DBAPI connection that can provide a cursor
-    dbapi_conn = conn.connection
-    with dbapi_conn.cursor() as cur:
-        s_buf = StringIO()
-        s_buf.writelines(
-            "\n".join(
-                map(lambda line: "|".join(map(str, map(parse_col, line))), data_iter)
-            )
-        )
-        s_buf.seek(0)
-
-        sql = "COPY {} ({}) FROM STDIN WITH null as 'null' DELIMITER '|'".format(
-            psycopg2.sql.Identifier(
-                *(table.schema, table.name) if table.schema else (table.name,)
-            ).as_string(cur),
-            psycopg2.sql.SQL(", ")
-            .join(
-                map(
-                    psycopg2.sql.Identifier,
-                    keys[1:] if keys and keys[0] == "index" else keys,
-                )
-            )
-            .as_string(cur),
-        )
-        try:
-            cur.copy_expert(sql=sql, file=s_buf)
-        except:
-            print(sql)
-            s_buf.seek(0)
-            pp(
-                dict(
-                    zip(
-                        keys[1:] if keys and keys[0] == "index" else keys,
-                        next(s_buf).split("|"),
-                    )
-                )
-            )
-            raise
+
+    with conn.connection.cursor() as cur:
+        table_name = psycopg2.sql.Identifier(
+            *(table.schema, table.name) if table.schema else (table.name,)
+        ).as_string(cur)
+
+    mgr = CopyManager(
+        conn.connection, table.name, keys[1:] if keys and keys[0] == "index" else keys
+    )
+    mgr.copy(data_iter)
+    conn.connection.commit()
 
 
 def csv_to_postgres_text(lines):

cdd_gae/tests/test_parquet_to_table.py

@@ -17,6 +17,7 @@
 import pyarrow as pa
 import pyarrow.parquet as pq
 import sqlalchemy
+from cdd.shared.pure_utils import rpartial
 from cdd.tests.utils_for_tests import unittest_main
 from psycopg2.extensions import AsIs, register_adapter
 from sqlalchemy import (
@@ -121,12 +122,15 @@ class TestParquetToTable(TestCase):
     )  # python: `pd.read_sql_query("SELECT * FROM").to_csv`
 
     def test_csv_to_postgres_text(self):
-        self.assertEqual(
-            self.copy_to_stdout_mock, csv_to_postgres_text(self.to_csv_mock)
+        self.assertListEqual(
+            *map(
+                rpartial(str.split, "\n"),
+                (self.copy_to_stdout_mock, csv_to_postgres_text(self.to_csv_mock)),
+            )
         )
 
     @unittest.skipUnless(
-        "RDBMS_URI" in environ, "RDMBS_URI env var must be set to for this test to run"
+        "RDBMS_URI" in environ, "RDMBS_URI env var must be set for this test to run"
     )
     def test_sqlalchemy_csv(self) -> None:
         """Reverse-engineer SQLalchemy's handling of complex types for cdd_gae's implementation"""
@@ -165,8 +169,9 @@ class Tbl:
                 session.commit()
                 self.assertEqual(Query([Tbl], session=session).count(), 3)
                 df = pd.read_sql_query("SELECT * FROM {}".format(table_name), engine)
+                # session.execute("TRUNCATE {}".format(table_name))
 
-                results = df.to_csv(index=False, sep="\t")
+                results = df.to_csv(index=False, sep="\t", header=False)
                 self.assertEqual(
                     results,
                     "timestamp_col\tjson_col\tarray_str_col\tarray_bigint_col\tarray_json_col\tid\n{0}".format(
@@ -185,7 +190,7 @@ class Tbl:
                 metadata.drop_all(tables=[table], bind=engine)
 
     @unittest.skipUnless(
-        "RDBMS_URI" in environ, "RDMBS_URI env var must be set to for this test to run"
+        "RDBMS_URI" in environ, "RDMBS_URI env var must be set for this test to run"
     )
     @with_own_table
     def test_parquet_to_table(self) -> None:

requirements.txt

@@ -1,10 +1,11 @@
 # Note: no dependencies are required on Python >=3.9; `black` is optional; `pyyaml` is mostly optional
-https://github.com/offscale/astor/archive/refs/heads/empty-annassign.zip
+#python-cdd
 black
+https://api.github.com/repos/offscale/cdd-python/zipball#egg=python_cdd
+https://github.com/offscale/astor/archive/refs/heads/empty-annassign.zip
+pgcopy
 psycopg
 pyarrow
-#python-cdd
-https://api.github.com/repos/offscale/cdd-python/zipball#egg=python_cdd
 pyyaml
 sqlalchemy
 typing-extensions