Repository: spotify/luigi
Issue category: Bug Report

[contrib.postgres] copy_from does not accept schema.table notation in most recent psycopg2 versions
<!---
We use GitHub issues mainly for tracking bugs and feature requests.
Questions for how to use luigi can be sent to the mailing list.

Currently, there are no strict procedures or guidelines for submitting issues.
In short, please just use common sense.

Common sense includes this at bare-minimum:

 * search for similar issues posted before creating a new issue.
 * Use markdown to format all code/logs. Issues which are hard to read
   when rendered on GitHub might be closed with a friendly reminder of this.
 * If applicable, reading relevant parts of the documentation.

Also, add steps to reproduce the bug, if applicable. Sample code would be nice too :)

For more information on how to submit valuable contributions,
see https://opensource.guide/how-to-contribute/#how-to-submit-a-contribution
-->

## Description

I'm trying to maintain an old (2018) project that includes a lot of Luigi tasks, amongst which there are some tasks derived from [`CopyToTable`](https://github.com/spotify/luigi/blob/master/luigi/contrib/postgres.py#L299).

In this project, the PostgreSQL database contains several schemas, and some data may be added to tables which are not in default `public` schema through `CopyToTable`-derived tasks.

However `CopyToTable` uses the `cursor.copy_from` method, which has been recently modified in `psycopg2` API (see *e.g* https://github.com/psycopg/psycopg2/issues/1294). Hence using Luigi with a recent `psycopg2` raises an error like `psycopg2.errors.UndefinedTable: relation "schema.table" does not exist`.

## Expected behavior

Taking into account the behavior change in psycopg2, considering `schema.table` notation for Postgres tables that are located in a dedicated schema.

## Minimal Working Example

Let's consider the following Python module (let's call it `luigi_copytotable.py`) :

```python
from luigi.contrib.postgres import CopyToTable
import pandas as pd


class SendToDB(CopyToTable):
    """Insert bike availability data into a PostgreSQL table                                  
    """

    host = "localhost"
    database = "my_db"
    user = "my_username"
    password = "my_password"

    columns = [('a', 'VARCHAR'), ('b', 'INT')]

    @property
    def table(self):
        return 'my_schema.my_table'

    def rows(self):
        df = pd.DataFrame({"a": ["foo", "bar", "wiz"], "b": [1, 2, 3]})
        for idx, row in df.iterrows():
            yield row.values
```

Running `luigi --local-scheduler --module luigi_copytotable SendToDB` throws:

```bash
16:04 $ luigi --local-scheduler --module luigi_copytotable SendToDB_
DEBUG: Checking if SendToDB() is complete
INFO: Informed scheduler that task   SendToDB__99914b932b   has status   PENDING
INFO: Done scheduling tasks
INFO: Running Worker with 1 processes
DEBUG: Asking scheduler for work...
DEBUG: Pending tasks: 1
INFO: [pid 494717] Worker Worker(salt=450412579, workers=1, host=*******, username=my_username, pid=494717) running   SendToDB_()
INFO: Done writing, importing at 2022-09-07 16:04:05.364381
INFO: Creating table my_schema.my_table
ERROR: [pid 494717] Worker Worker(salt=450412579, workers=1, host=*******, username=my_username, pid=494717) failed    SendToDB()
Traceback (most recent call last):
  File "/home/rdelhome/.virtualenvs/jitenv/lib/python3.10/site-packages/luigi/worker.py", line 198, in run
    new_deps = self._run_get_new_deps()
  File "/home/rdelhome/.virtualenvs/jitenv/lib/python3.10/site-packages/luigi/worker.py", line 138, in _run_get_new_deps
    task_gen = self.task.run()
  File "/home/rdelhome/.virtualenvs/jitenv/lib/python3.10/site-packages/luigi/contrib/postgres.py", line 403, in run
    self.copy(cursor, tmp_file)
  File "/home/rdelhome/.virtualenvs/jitenv/lib/python3.10/site-packages/luigi/contrib/postgres.py", line 358, in copy
    cursor.copy_from(
psycopg2.errors.UndefinedTable: relation "my_schema.my_table" does not exist

DEBUG: 1 running tasks, waiting for next task to finish
INFO: Informed scheduler that task   SendToDB__99914b932b   has status   FAILED
DEBUG: Asking scheduler for work...
DEBUG: Done
DEBUG: There are no more tasks to run at this time
DEBUG: There are 1 pending tasks possibly being run by other workers
DEBUG: There are 1 pending tasks unique to this worker
DEBUG: There are 1 pending tasks last scheduled by this worker
INFO: Worker Worker(salt=450412579, workers=1, host=*********, username=my_username, pid=494717) was stopped. Shutting down Keep-Alive thread
INFO: 
===== Luigi Execution Summary =====

Scheduled 1 tasks of which:
* 1 failed:
    - 1 SendToDB()

This progress looks :( because there were failed tasks

===== Luigi Execution Summary =====
```

## Hints for resolution

As suggested in the psycopg2 issue, use `copy_expert` ? Or maybe modify the `if` predicate in https://github.com/spotify/luigi/blob/master/luigi/contrib/postgres.py#L357, to choose the `else` case if `copy_from` is not happy...

(Note: As a temporary solution, I've downgraded my  `psycopg2` version to `<2.9` to make it work.)

## Related issue

See on psycopg2 project: https://github.com/psycopg/psycopg2/issues/1294



Functions to edit:
luigi/contrib/postgres.py:CopyToTable.copy

Patch:
diff --git a/luigi/contrib/postgres.py b/luigi/contrib/postgres.py
index 719b80a4d7..19e96e8180 100644
--- a/luigi/contrib/postgres.py
+++ b/luigi/contrib/postgres.py
@@ -356,16 +356,15 @@ def copy(self, cursor, file):
         else:
             raise Exception('columns must consist of column strings or (column string, type string) tuples (was %r ...)' % (self.columns[0],))
 
-        # cursor.copy_from is not available in pg8000
-        if hasattr(cursor, 'copy_from'):
-            cursor.copy_from(
-                file, self.table, null=r'\\N', sep=self.column_separator, columns=column_names)
+        copy_sql = (
+            "COPY {table} ({column_list}) FROM STDIN "
+            "WITH (FORMAT text, NULL '{null_string}', DELIMITER '{delimiter}')"
+        ).format(table=self.table, delimiter=self.column_separator, null_string=r'\\N',
+                 column_list=", ".join(column_names))
+        # cursor.copy_expert is not available in pg8000
+        if hasattr(cursor, 'copy_expert'):
+            cursor.copy_expert(copy_sql, file)
         else:
-            copy_sql = (
-                "COPY {table} ({column_list}) FROM STDIN "
-                "WITH (FORMAT text, NULL '{null_string}', DELIMITER '{delimiter}')"
-            ).format(table=self.table, delimiter=self.column_separator, null_string=r'\\N',
-                     column_list=", ".join(column_names))
             cursor.execute(copy_sql, stream=file)
 
     def run(self):