Shakshi3104 commited on
Commit
133e909
1 Parent(s): e319108

[add] Implement notion_db to load contents from Notion database

Browse files
Files changed (1) hide show
  1. model/data/notion_db.py +216 -0
model/data/notion_db.py ADDED
@@ -0,0 +1,216 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from pathlib import Path
3
+ import abc
4
+
5
+ import pandas as pd
6
+ from dotenv import load_dotenv
7
+
8
+ import notion_client as nt
9
+ from notion2md.exporter.block import StringExporter
10
+
11
+ from loguru import logger
12
+
13
+
14
+ class BaseNotionDatabase:
15
+ """
16
+ Notion DBからページのコンテンツを取り出すベースのクラス
17
+ """
18
+ def __init__(self):
19
+ load_dotenv()
20
+ self.notion_database_id = os.getenv("NOTION_DATABASE_ID")
21
+ self.integration_token = os.getenv("INTEGRATION_TOKEN")
22
+
23
+ # notion2mdの環境変数
24
+ os.environ["NOTION_TOKEN"] = os.getenv("INTEGRATION_TOKEN")
25
+
26
+ self.notion_client = nt.Client(auth=self.integration_token)
27
+
28
+ def load_database(self) -> list[dict]:
29
+ """
30
+ Notion DBのページ一覧を取得
31
+
32
+ Returns:
33
+
34
+ """
35
+ results = []
36
+ has_more = True
37
+ start_cursor = None
38
+
39
+ while has_more:
40
+ db = self.notion_client.databases.query(
41
+ **{
42
+ "database_id": self.notion_database_id,
43
+ "start_cursor": start_cursor
44
+ }
45
+ )
46
+ # 100件までしか1回に取得できない
47
+ # 100件以上ある場合 has_more = True
48
+ has_more = db["has_more"]
49
+ # 次のカーソル
50
+ start_cursor = db["next_cursor"]
51
+
52
+ # 取得結果
53
+ results += db["results"]
54
+
55
+ return results
56
+
57
+ @abc.abstractmethod
58
+ def load_content(self) -> list[dict]:
59
+ """
60
+ Notion DBのページの中身をdictで返す
61
+ Returns:
62
+
63
+ """
64
+ raise NotImplementedError
65
+
66
+
67
+ class SakurapDB(BaseNotionDatabase):
68
+ def load_database(self) -> list[dict]:
69
+ """
70
+ Notion DBのページ一覧を取得
71
+
72
+ Returns:
73
+ results:
74
+ list[dict]
75
+
76
+ """
77
+ results = []
78
+ has_more = True
79
+ start_cursor = None
80
+
81
+ while has_more:
82
+ # "Rap詞 : 櫻井翔"がTrueのもののみ取得
83
+ db = self.notion_client.databases.query(
84
+ **{
85
+ "database_id": self.notion_database_id,
86
+ "filter": {
87
+ "property": "Rap詞 : 櫻井翔",
88
+ "checkbox": {
89
+ "equals": True
90
+ }
91
+ },
92
+ "start_cursor": start_cursor
93
+ }
94
+ )
95
+ # 100件までしか1回に取得できない
96
+ # 100件以上ある場合 has_more = True
97
+ has_more = db["has_more"]
98
+ # 次のカーソル
99
+ start_cursor = db["next_cursor"]
100
+
101
+ # 取得結果
102
+ results += db["results"]
103
+
104
+ return results
105
+
106
+ def __load_blocks(self, block_id: str) -> str:
107
+ """
108
+ Notionのページをプレーンテキストで取得する (Notion Official API)
109
+
110
+ Parameters
111
+ ----------
112
+ block_id:
113
+ str, Block ID
114
+
115
+ Returns
116
+ -------
117
+ texts:
118
+ str
119
+ """
120
+ block = self.notion_client.blocks.children.list(
121
+ **{
122
+ "block_id": block_id
123
+ }
124
+ )
125
+
126
+ # プレーンテキストを繋げる
127
+ def join_plain_texts():
128
+ text = [blck["paragraph"]["rich_text"][0]["plain_text"] if len(blck["paragraph"]["rich_text"])
129
+ else "\n" for blck in block["results"]]
130
+
131
+ texts = "\n".join(text)
132
+ return texts
133
+
134
+ return join_plain_texts()
135
+
136
+ def load_content(self) -> list[dict]:
137
+ """
138
+ Notion DBのページの中身をdictで返す
139
+
140
+ Returns:
141
+ lyrics:
142
+ list[dict]
143
+ """
144
+
145
+ # DBのページ一覧を取得
146
+ db_results = self.load_database()
147
+ logger.info("🚦 [Notion] load database...")
148
+
149
+ # コンテンツ一覧
150
+ lyrics = []
151
+
152
+ logger.info("🚦 [Notion] start to load each page content ...")
153
+ # 各ページの処理
154
+ for result in db_results:
155
+ block_id = result["id"]
156
+ # rap_lyric = self.__load_blocks(block_id)
157
+
158
+ # Markdown形式でページを取得
159
+ rap_lyric = StringExporter(block_id=block_id).export()
160
+ # Markdownの修飾子を削除
161
+ rap_lyric = rap_lyric.replace("\n\n", "\n").replace("<br/>", "\n").replace("*", "")
162
+
163
+ lyrics.append(
164
+ {
165
+ "title": result["properties"]["名前"]["title"][0]["plain_text"],
166
+ "content": rap_lyric
167
+ }
168
+ )
169
+
170
+ logger.info("🚦 [Notion] Finish to load.")
171
+
172
+ return lyrics
173
+
174
+
175
+ def fetch_sakurap_corpus(filepath: str, refetch=False) -> pd.DataFrame:
176
+ """
177
+ サクラップのコーパスを取得する
178
+ CSVファイルが存在しないときにNotionから取得する
179
+
180
+ Parameters
181
+ ----------
182
+ filepath:
183
+ str
184
+ refetch:
185
+ bool
186
+
187
+ Returns
188
+ -------
189
+
190
+ """
191
+ filepath = Path(filepath)
192
+
193
+ if not filepath.exists() or refetch:
194
+ # CSVファイルを保存するディレクトリが存在しなかったら作成する
195
+ if not filepath.parent.exists():
196
+ logger.info(f"🚦 [Notion] mkdir {str(filepath.parent)} ...")
197
+ filepath.parent.mkdir(parents=True, exist_ok=True)
198
+
199
+ logger.info("🚦 [Notion] fetch from Notion DB ...")
200
+ # dictを取得
201
+ rap_db = SakurapDB()
202
+ lyrics = rap_db.load_content()
203
+
204
+ lyrics_df = pd.DataFrame(lyrics)
205
+ lyrics_df.to_csv(filepath, index=False)
206
+ else:
207
+ logger.info("🚦 [Notion] load CSV file.")
208
+
209
+ lyrics_df = pd.read_csv(filepath)
210
+
211
+ return lyrics_df
212
+
213
+
214
+ if __name__ == "__main__":
215
+ sakurap_db = SakurapDB()
216
+ lyrics = sakurap_db.load_content()