-
Notifications
You must be signed in to change notification settings - Fork 0
/
reddit_to_db.py
92 lines (78 loc) · 3.43 KB
/
reddit_to_db.py
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
import asyncpraw
import asyncio
import psycopg
import time
from private import reddit_details, db_details
from aiostream import stream
import datetime
def parse_item(item: asyncpraw.reddit.Comment | asyncpraw.reddit.Submission) -> dict:
# the item can be either a comment or a submission
# comments have body attr, submissions have title and selftext attrs
# lets gather all of them and set it to None if it does not exist
attrs = [
'id',
'created_utc',
'title', # only submissions
'body', # only comments
'selftext', # only submissions
]
return {key: getattr(item, key, None) for key in attrs}
async def reddit_stream(reddit_details, subreddit):
# Does the following:
# 1. create comment stream
# 2. create submission stream
# 3. merge both streams into a single stream
# 4. yield stream items
async def comment_stream(subreddit: asyncpraw.reddit.Subreddit) -> asyncpraw.Reddit.comment:
async for comment in subreddit.stream.comments():
yield comment
async def submission_stream(subreddit: asyncpraw.reddit.Subreddit) -> asyncpraw.Reddit.submission:
async for submission in subreddit.stream.submissions():
yield submission
reddit = asyncpraw.Reddit(**reddit_details)
async with reddit:
subreddit_ = await reddit.subreddit(subreddit)
merged_stream = stream.merge(comment_stream(subreddit_), submission_stream(subreddit_))
async with merged_stream.stream() as streamer:
async for item in streamer:
yield item
async def main(reddit_details: dict, subreddit: str = 'cryptocurrency') -> None:
# define db query
columns = ['_id', 'created_utc', 'title', 'body', 'selftext']
query = f"""
INSERT INTO {subreddit} ({', '.join(columns)}) VALUES (%s, %s, %s, %s, %s)
ON CONFLICT (_id)
DO UPDATE SET ({', '.join(columns[1:])})
= ({", ".join(['EXCLUDED.' + col for col in columns[1:]])})
"""
count = 0
# it is recommended to use global connection_pool object
# and connect to the db once and push transactions through it
# but cant install the required dependencies.
# So lets just wrap everything around a single connection for now.
# TODO: https://www.psycopg.org/psycopg3/docs/advanced/pool.html
async with await psycopg.AsyncConnection.connect(**db_details, dbname='reddit') as aconn:
async with aconn:
streamer = reddit_stream(reddit_details, subreddit)
async for item in streamer:
if item:
item_ = parse_item(item)
await aconn.cursor().execute(query, tuple(item_.values()))
await aconn.commit()
count += 1
if count % 10000 == 0 or count == 1:
print(
datetime.datetime.now(datetime.timezone.utc).strftime('%Y-%m-%d %H:%M:%S'),
'comms/subms pushed to db:', count
)
if __name__ == "__main__":
while True:
try:
asyncio.run(main(reddit_details=reddit_details, subreddit='cryptocurrency'))
except Exception as e:
print(
'Failed to run the main loop. Time:',
datetime.datetime.now(datetime.timezone.utc).strftime('%Y-%m-%d %H:%M:%S'),
'Error: ', e
)
time.sleep(10)