You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
2 years ago
|
import asyncio
|
||
|
|
import asyncpg
|
||
|
|
import config
|
||
|
|
import os
|
||
|
|
|
||
|
|
conn = None
|
||
|
|
|
||
|
|
class Install:
|
||
|
|
dbi: asyncpg.Connection
|
||
|
|
|
||
|
|
async def run(self):
|
||
|
|
self.dbi = await asyncpg.connect(**config.DATABASE)
|
||
|
|
args = os.sys.argv
|
||
|
|
if "--force" in args:
|
||
|
|
await self.drop_table()
|
||
|
|
|
||
|
|
await self.create_table()
|
||
|
|
|
||
|
|
async def drop_table(self):
|
||
|
|
await self.dbi.execute("DROP TABLE IF EXISTS embedding_search_title_index;")
|
||
|
|
print("Table dropped")
|
||
|
|
|
||
|
|
async def create_table(self):
|
||
|
|
await self.dbi.execute("""
|
||
|
|
CREATE TABLE embedding_search_title_index (
|
||
|
|
id SERIAL PRIMARY KEY,
|
||
|
|
sha1 VARCHAR(40) NOT NULL UNIQUE,
|
||
|
|
title TEXT NOT NULL,
|
||
|
|
rev_id INT8 NOT NULL,
|
||
|
|
embedding VECTOR(%d) NOT NULL
|
||
|
|
);
|
||
|
|
""" % (config.EMBEDDING_VECTOR_SIZE))
|
||
|
|
await self.dbi.execute("CREATE INDEX embedding_search_title_index_embedding_idx ON embedding_search_title_index USING ivfflat (embedding vector_cosine_ops);")
|
||
|
|
print("Table created")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
install = Install()
|
||
|
|
|
||
|
|
loop = asyncio.get_event_loop()
|
||
|
|
loop.run_until_complete(install.run())
|