-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcache_back.py
More file actions
236 lines (202 loc) Β· 8.1 KB
/
cache_back.py
File metadata and controls
236 lines (202 loc) Β· 8.1 KB
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
import psycopg2
from sqlalchemy import create_engine
import pandas as pd
import headers
import hashlib
import dependency as dep
import os
DB_NAME = ""
DB_USER = ""
DB_PASS = ""
DB_HOST = ""
DB_PORT = 5432 # default
BLOB_TABLE_NAME = "pipeline_blobs"
NOTEBOOK_NAME = ""
NOTEBOOK_CODE = ""
cached_objects = {}
cache_outputs = {}
CACHE_TABLE_NAME = "cached_tables"
def init_session(db_name, db_user, db_pass, db_host, db_port=5432, notebook_name=''):
""" Initialises Database parameters for connection to postgres """
global DB_NAME, DB_USER, DB_PASS, DB_HOST, DB_PORT, NOTEBOOK_CODE, NOTEBOOK_NAME
DB_NAME = db_name
DB_USER = db_user
DB_PASS = db_pass
DB_HOST = db_host
DB_PORT = db_port
NOTEBOOK_NAME = notebook_name
# currently this is causing problem, so temporarily disabled
#dep.handle_imports(notebook_name + '.ipynb')
print(DB_USER,DB_PASS, DB_HOST, DB_PORT)
print('init session invoked')
def _connect():
""" Connect to postgres and retrieve all cached variables back into the script """
conn = psycopg2.connect(database=DB_NAME,
user=DB_USER,
password=DB_PASS,
host=DB_HOST,
port=DB_PORT)
print('_connect invoked')
return conn
def insert(df, destination_db_table):
""" Inserts a given dataframe into postgres """
try:
conn_string = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
print(conn_string)
db = create_engine(conn_string)
print('created engine')
conn = db.connect()
print('connected')
df.to_sql(destination_db_table, con=conn, if_exists='replace', index=False, chunksize = 1000)
print(f"Cached Dataframe successfully to table: {destination_db_table}")
except Exception as e:
print(e.args[0])
# When executing code,
# 1. get all lists of pandas dataframes
# 2. insert them into cached table (cached_tables)
# the default behaviour of commit should be every pandas table needs to be cached
def execute_as_plpython(notebook_path, function_name):
""" Takes a jupyter notebook and runs it as a plpython function on Postgres Server """
try:
plpython_query = headers.generate_query(notebook_path,
function_name,
is_query=False)
conn = psycopg2.connect(database=DB_NAME,
user=DB_USER,
password=DB_PASS,
host=DB_HOST,
port=DB_PORT)
print(f"EXECUTING -->\n{plpython_query}")
cur = conn.cursor()
cur.execute(plpython_query)
print('RUNNING SCRIPT')
cur.execute(f"SELECT {function_name}();")
res = cur.fetchone()
print("these variables were cached:")
# the cursor returns the cache_outputs from plpython instance as "string" format.
# therefore, some code for properly parsing is needed. :)
global cache_outputs
for i in res:
cache_outputs = eval(i)
for i in cache_outputs.keys():
print(i)
print('Successful execution of plpython')
conn.commit()
cur.close()
except Exception as e:
print(e.args[0])
def add_to_cache(table: pd.DataFrame, object: str):
'''
please make the table varialbe and the object name the same
i.e., add_to_cache(df1, 'df1')
'''
cached_objects[object] = table
print(f"Object added to cache. Current Cache: {cached_objects.keys()}")
def view_cache():
print(cached_objects)
def remove_from_cache(object_name):
if object_name in cached_objects.keys():
del cached_objects[object_name]
print(f"Removed {object_name}")
else:
raise Exception(f"{object_name} not found.")
def send_blob(notebook_path, file_name):
_create_blob_table()
try:
conn = _connect()
cur = conn.cursor()
file_data = read_notebook_as_binary(notebook_path)
global NOTEBOOK_CODE
NOTEBOOK_CODE = headers.get_code_from_notebook(filename=notebook_path)
function_name = headers.get_notebook_name(notebook_path) + '_script'
plpython_script = headers.generate_query(notebook_path, function_name, is_query=True)
print('generated query')
print(plpython_script)
blob = psycopg2.Binary(file_data)
print('file read as binary')
query = f"INSERT INTO {BLOB_TABLE_NAME} (file_name, source_notebook, plscript) VALUES('{file_name}',{blob},'''{plpython_script}''')"
print(query)
cur.execute(query)
print('Blob inserted')
conn.commit()
cur.close()
execute_as_plpython(notebook_path, 'execute_plpython')
print("this is the rewritten pipeline")
print("-------------------------------")
except Exception as e:
print(e.args[0])
def _create_blob_table():
try:
conn = _connect()
cur = conn.cursor()
query = f"CREATE TABLE IF NOT EXISTS {BLOB_TABLE_NAME} (id SERIAL PRIMARY KEY, upload_date TIMESTAMP default current_timestamp, file_name TEXT, source_notebook BYTEA, plscript TEXT, updated_notebook BYTEA);"
print(query)
cur.execute(query)
conn.commit()
cur.close()
except Exception as e:
print(e.args)
def read_notebook_as_binary(notebook_path):
with open(notebook_path, 'rb') as file:
data = file.read()
return data
def cache_from_list():
try:
print('Caching Objects ... ')
for df_name, df in cached_objects.items():
print(f"inserting {df_name} ...")
df_table = generate_var_name(df_name)
insert(df,destination_db_table=df_table)
cache_outputs[df_name] = f"SELECT * FROM {df_table}"
# rewrite_pipeline()
except Exception as e:
print(e.args[0])
def read_existing_cache(notebook_path: str):
"""
Attempts to check whether DBMS already has existing cache for current notebook.
If it does not exist, prints that the cache does not exist.
"""
print("Attemping to retrieve cache (if existing)")
try:
conn = _connect()
cur = conn.cursor()
file_data = read_notebook_as_binary(notebook_path)
blob = psycopg2.Binary(file_data)
query = f"SELECT * FROM {BLOB_TABLE_NAME} WHERE source_notebook = {blob}"
cur = conn.cursor()
cur.execute(query)
res = cur.fetchall()
print(res)
cur.close()
except Exception as e:
print(e.args)
def generate_var_name(df_name, filename=NOTEBOOK_NAME):
hash = hashlib.md5(filename.encode()).hexdigest()
generated_name = df_name + '_' + hash
return generated_name
def rewrite_pipeline():
global NOTEBOOK_CODE
updated_code = ''
for df_name, query in cache_outputs.items():
# commenting and also adding newly variable assignments
updated_code = headers.comment_line_by_var_usage(df_name, query, NOTEBOOK_CODE)
NOTEBOOK_CODE = updated_code
#headers.rewrite_var_definition(df_name, query, updated_code)
# handling sql connection code in front of everything
NOTEBOOK_CODE = get_sql_conn_code() + updated_code
print('*********************************')
print(NOTEBOOK_CODE)
def get_sql_conn_code():
explanation = "# FOLLOWING IMPORT GENERATED BY CACHEBACK"
import_code = "import psycopg2\nimport pandas as pd"
l1 = f"cacheback_conn = psycopg2.connect(database='{DB_NAME}', user='{DB_USER}', password='{DB_PASS}', host='{DB_HOST}', port={DB_PORT})"
l2 = "cacheback_cur = cacheback_conn.cursor()"
endline = "#------------------------------------------"
return explanation + '\n' + import_code + '\n' + l1 + '\n' + l2 + '\n' + endline + '\n\n'
def return_new_notebook(nb_name):
# first save rewritten notebook code as .py
file = open(f"./{nb_name}_rewritten.py", 'w')
file.write(NOTEBOOK_CODE)
file.close()
# then convert .py to .ipynb using jupytext
os.system(f'jupytext --to notebook {nb_name}_rewritten.py')