PoolToolBackend-DB/pg_utils.py at master · PoolTool-io/PoolToolBackend-DB · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
from config import *
import psycopg2.extras
from psycopg2.extras import Json
from functools import wraps
import time


def retry(fn=None):
    @wraps(fn)
    def wrapper(*args, **kw):
        cls = args[0]
        while True:
            try:
                return fn(*args, **kw)
            except (psycopg2.InterfaceError, psycopg2.OperationalError) as e:
                print(e)
                print ("\nDatabase Connection [InterfaceError or OperationalError]")
                print ("Idle for %s seconds" % (2))
                time.sleep(2)
                cls._connect()
    return wrapper

class pg_utils:

    def __init__(self,connname="testdbconn"):
        self.conn = None
        self.cur1 = None
        self.cur2 = None
        self.cur3 = None
        self.connname=connname

        self._connect()

    def _connect(self):
        self.conn = psycopg2.connect(application_name=self.connname, database=database, user=user, password=password, host=host, port=port)
        self.cur1 = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        self.cur2 = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        self.cur3 = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)

    def __del__(self):
        # Maybe there is a connection but no cursor, whatever close silently!
        try:
            self.conn.close()
        except:
            pass
        self.conn = None

    @retry
    def cur1_create(self):
        if self.cur1 == None or self.cur1.closed:
            self.cur1 = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        return True

    @retry
    def cur2_create(self):
        if self.cur2 == None or self.cur2.closed:
            self.cur2 = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        return True

    @retry
    def cur3_create(self):
        if self.cur3 == None or self.cur3.closed:
            self.cur3 = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        return True
    @retry
    def create_server_side_cursor(self, name, itersize=10000):
        cursor = self.conn.cursor(name=name, cursor_factory=psycopg2.extras.DictCursor)
        cursor.itersize = itersize
        return cursor

    def cur1_close(self):
        self.cur1.close()
        self.cur1=None
        return True

    def cur2_close(self):
        self.cur2.close()
        self.cur2=None
        return True

    def cur3_close(self):
        self.cur3.close()
        self.cur3=None
        return True

    @retry
    def cur1_execute(self,sql,params=None): # pass here required params to get data from DB
        return self.cur1.execute(sql,params)

    @retry
    def cur1_fetchone(self): # pass here required params to get data from DB
        return self.cur1.fetchone()

    @retry
    def cur1_fetchall(self): # pass here required params to get data from DB
        return self.cur1.fetchall()

    @retry
    def cur1_fetchmany(self, size=10000): # Fetch a batch of rows
        return self.cur1.fetchmany(size)
    @retry
    def cur1_copy_expert(self, sql, output_file): # Use the COPY command
        with open(output_file, 'w') as f:
            self.cur1.copy_expert(sql, f)
    @retry
    def cur1_copy_expert_with_headers(self, sql, output_file, headers):
        with open(output_file, 'w') as f:
            f.write(headers)  # Write the custom headers to the CSV file
            self.cur1.copy_expert(sql, f)
    @retry
    def cur2_execute(self,sql,params=None): # pass here required params to get data from DB
        return self.cur2.execute(sql,params)

    @retry
    def cur2_fetchone(self): # pass here required params to get data from DB
        return self.cur2.fetchone()

    @retry
    def cur3_execute(self,sql,params=None): # pass here required params to get data from DB
        return self.cur3.execute(sql,params)

    @retry
    def cur3_fetchone(self): # pass here required params to get data from DB
        return self.cur3.fetchone()

    @retry
    def server_side_execute(self, cursor_name, sql, params=None, itersize=10000):
        cursor = self.create_server_side_cursor(cursor_name, itersize)
        cursor.execute(sql, params)
        return cursor  # Return the cursor for iteration

    @retry
    def conn_commit(self):
        return self.conn.commit()