fix all the psycopg2 refs

segasai · Apr 24, 2024 · 227a56d · 227a56d
1 parent 1f71f36
commit 227a56d
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 18 deletions.
diff --git a/py/sqlutilpy/sqlutil.py b/py/sqlutilpy/sqlutil.py
@@ -8,7 +8,6 @@
 import collections
 import warnings
 from numpy.core import numeric as sb
-from select import select
 
 try:
     import astropy.table as atpy
@@ -21,7 +20,6 @@
     # pandas is not installed
     pandas = None
 
-from io import BytesIO as StringIO
 import queue
 
 _WAIT_SELECT_TIMEOUT = 10
@@ -54,7 +52,7 @@ def getConnection(db=None,
         The name of the database (in case of PostgreSQL) or filename in
         case of sqlite db
     driver :  string
-        The db driver (either 'psycopg2' or 'sqlite3')
+        The db driver (either 'psycopg' or 'sqlite3')
     user : string, optional
         Username
     password: string, optional
@@ -73,6 +71,10 @@ def getConnection(db=None,
 
     """
     if driver == 'psycopg2':
+        warnings.warn(
+            'psycopg2 driver is not supported anymore using psycopg instead')
+        driver = 'psycopg'
+    if driver == 'psycopg':
         conn_dict = dict()
         if db is not None:
             conn_dict['dbname'] = db
@@ -100,6 +102,10 @@ def getCursor(conn, driver=None, preamb=None, notNamed=False):
     Retrieve the database cursor
     """
     if driver == 'psycopg2':
+        warnings.warn(
+            'psycopg2 driver is not supported anymore using psycopg instead')
+        driver = 'psycopg'
+    if driver == 'psycopg':
         cur = conn.cursor()
         if preamb is not None:
             cur.execute(preamb)
@@ -241,7 +247,7 @@ def __getDType(row, typeCodes, strLength):
 def get(query,
         params=None,
         db="wsdb",
-        driver="psycopg2",
+        driver="psycopg",
         user=None,
         password=None,
         host=None,
@@ -278,7 +284,7 @@ def get(query,
     db : string
         The name of the database
     driver : string, optional
-        The sql driver to be used (psycopg2 or sqlite3)
+        The sql driver to be used (psycopg or sqlite3)
     user : string, optional
         User name for the DB connection
     password : string, optional
@@ -333,6 +339,11 @@ def get(query,
         proc = None
         colNames = []
         if driver == 'psycopg2':
+            warnings.warn(
+                'psycopg2 driver is not supported anymore using psycopg instead'
+            )
+            driver = 'psycopg'
+        if driver == 'psycopg':
             try:
                 while (True):
                     # Iterating over the cursor, retrieving batches of results
@@ -447,7 +458,7 @@ def get(query,
 def execute(query,
             params=None,
             db='wsdb',
-            driver="psycopg2",
+            driver="psycopg",
             user=None,
             password=None,
             host=None,
@@ -467,7 +478,7 @@ def execute(query,
     db : string
         Database name
     driver : string
-        Driver for the DB connection ('psycopg2' or 'sqlite3')
+        Driver for the DB connection ('psycopg' or 'sqlite3')
     user : string, optional
         user name for the DB connection
     password : string, optional
@@ -600,7 +611,7 @@ def upload(tableName,
            arrays,
            names=None,
            db="wsdb",
-           driver="psycopg2",
+           driver="psycopg",
            user=None,
            password=None,
            host=None,
@@ -627,7 +638,7 @@ def upload(tableName,
     db: string
          Databas name
     driver: string
-         Python database driver "psycopg2",
+         Python database driver "psycopg",
     user: string,
     password: string
     host: string
@@ -746,7 +757,7 @@ def local_join(query,
                arrays,
                names,
                db=None,
-               driver="psycopg2",
+               driver="psycopg",
                user=None,
                password=None,
                host=None,

diff --git a/tests/test_pg.py b/tests/test_pg.py
@@ -1,7 +1,7 @@
 import pytest
 import string
 import os
-import psycopg2
+import psycopg
 import numpy as np
 import sqlutilpy as sqlutil
 import time
@@ -37,8 +37,8 @@ def getconn():
 
 @pytest.fixture
 def setup():
-    conn = psycopg2.connect('dbname=%s user=%s host=%s password=%s' %
-                            (PG_DB, PG_USER, PG_HOST, PG_PASS))
+    conn = psycopg.connect('dbname=%s user=%s host=%s password=%s' %
+                           (PG_DB, PG_USER, PG_HOST, PG_PASS))
     cur = conn.cursor()
     cur.execute('''
     create unlogged table sqlutil_test (sicol smallint, intcol int,

diff --git a/tests/test_sqlite.py b/tests/test_sqlite.py
@@ -1,12 +1,7 @@
 import pytest
-import string
 import os
-import psycopg2
-import numpy as np
 import sqlutilpy as sqlutil
 import sqlite3
-import time
-import killer
 
 
 @pytest.fixture()