@@ -47,6 +47,10 @@ def migrate_csv(
47
47
log .info (f"[yellow]Loading: { filename } " )
48
48
df = pd .read_csv (filename , header = None , skiprows = skip_rows , names = names )
49
49
log .info (f"Size of dataframe: { len (df )} " )
50
+ if model .__name__ == "Referring" : # Add date since this isn't included
51
+ file_date = filename .name [:10 ]
52
+ df .insert (loc = 0 , column = "date" , value = file_date )
53
+
50
54
if model .__name__ == "Paths" :
51
55
repository_names = [a .split ("/" )[2 ] for a in df ["path" ].values ]
52
56
df .insert (1 , "repository_name" , repository_names )
@@ -56,7 +60,12 @@ def migrate_csv(
56
60
query_results = list (
57
61
map (func , df ["repository_name" ], df ["date" ], df ["path" ])
58
62
)
59
- else :
63
+ elif model .__name__ == "Referring" :
64
+ func = partial (query_referring , engine = engine , model = model )
65
+ query_results = list (
66
+ map (func , df ["repository_name" ], df ["date" ], df ["site" ])
67
+ )
68
+ else : # For Clone and Traffic
60
69
func = partial (query , engine = engine , model = model )
61
70
query_results = list (map (func , df ["repository_name" ], df ["date" ]))
62
71
@@ -125,3 +134,25 @@ def query_path(
125
134
return result .one ()
126
135
except NoResultFound :
127
136
return
137
+
138
+
139
+ def query_referring (
140
+ repository_name : str ,
141
+ date : str ,
142
+ site : str ,
143
+ engine : Engine ,
144
+ model : Union [Type [SQLModel ], Referring ],
145
+ ) -> Union [SQLModel , Referring , None ]:
146
+
147
+ with Session (engine ) as session :
148
+ result = session .exec (
149
+ select (model ).where (
150
+ model .repository_name == repository_name ,
151
+ model .date == date ,
152
+ model .site == site ,
153
+ )
154
+ )
155
+ try :
156
+ return result .one ()
157
+ except NoResultFound :
158
+ return
0 commit comments