-
-
Notifications
You must be signed in to change notification settings - Fork 22
/
pgvector_util.py
84 lines (80 loc) · 2.59 KB
/
pgvector_util.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
from vdf_io.util import set_arg_from_input, set_arg_from_password
def make_pgv_parser(DB_NAME_SLUG, subparsers):
parser_pgvector = subparsers.add_parser(
DB_NAME_SLUG, help="Import data to pgvector"
)
parser_pgvector.add_argument(
"--platform",
type=str,
choices=["supabase", "neon", "tembo", "any"],
default="any",
help="Platform to connect to",
)
parser_pgvector.add_argument(
"--schema",
type=str,
help="Schema of the Postgres instance (default: public)",
default="public",
)
parser_pgvector.add_argument(
"--connection_string",
type=str,
help="Connection string to Postgres instance",
)
parser_pgvector.add_argument("--user", type=str, help="User of Postgres instance")
parser_pgvector.add_argument(
"--password", type=str, help="Password of Postgres instance"
)
parser_pgvector.add_argument("--host", type=str, help="Host of Postgres instance")
parser_pgvector.add_argument("--port", type=str, help="Port of Postgres instance")
parser_pgvector.add_argument(
"--dbname", type=str, help="Database name of Postgres instance"
)
parser_pgvector.add_argument(
"--tables", type=str, help="Postgres tables to export (comma-separated)"
)
return parser_pgvector
def set_pgv_args_from_prompt(args):
set_arg_from_input(
args,
"connection_string",
"Enter the connection string to Postgres instance: ",
str,
)
if not args.get("connection_string"):
set_arg_from_input(
args,
"user",
"Enter the user of Postgres instance (default: postgres): ",
str,
default_value="postgres",
)
set_arg_from_password(
args,
"password",
"Enter the password of Postgres instance (default: postgres): ",
)
if not args.get("password"):
# If password is not provided, set it to "postgres"
args["password"] = "postgres"
set_arg_from_input(
args,
"host",
"Enter the host of Postgres instance: ",
str,
default_value="localhost",
)
set_arg_from_input(
args,
"port",
"Enter the port of Postgres instance: ",
str,
default_value="5432",
)
set_arg_from_input(
args,
"dbname",
"Enter the database name of Postgres instance: ",
str,
default_value="postgres",
)