remote/sqlite-utils
1
0
Fork 0
mirror of https://github.com/simonw/sqlite-utils.git synced 2026-07-23 01:14:31 +02:00
Code Issues Projects Releases Packages Wiki Activity

.get() method plus support for compound primary keys (#40)

Browse source 
* create_table now handles compound primary keys
* CLI now accepts multiple --pk for compound primary keys
* Docs for compound primary keys with CLI and Python library
* New .get() method plus documentation

Closes #36, closes #39
This commit is contained in:
Simon Willison 2019-07-14 21:28:51 -07:00 • committed by GitHub
commit c65b67ca46
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 173 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

6
sqlite_utils/cli.py
View file

@ -302,7 +302,9 @@ def insert_upsert_options(fn):
),
click.argument("table"),
click.argument("json_file", type=click.File(), required=True),
click.option("--pk", help="Column to use as the primary key, e.g. id"),
click.option(
"--pk", help="Columns to use as the primary key, e.g. id", multiple=True
),
click.option("--nl", is_flag=True, help="Expect newline-delimited JSON"),
click.option("-c", "--csv", is_flag=True, help="Expect CSV"),
click.option(
@ -348,6 +350,8 @@ def insert_upsert_implementation(
if nl and csv:
click.echo("Use just one of --nl and --csv", err=True)
return
if pk and len(pk) == 1:
pk = pk[0]
if csv:
reader = csv_std.reader(json_file)
headers = next(reader)

67
sqlite_utils/db.py
View file

@ -79,6 +79,10 @@ class BadPrimaryKey(Exception):
pass
class NotFoundError(Exception):
pass
class Database:
def __init__(self, filename_or_conn):
if isinstance(filename_or_conn, str):
@ -205,11 +209,15 @@ class Database:
raise AlterError(
"No such column: {}.{}".format(fk.other_table, fk.other_column)
)
extra = ""
column_defs = []
# ensure pk is a tuple
single_pk = None
if isinstance(pk, str):
single_pk = pk
for column_name, column_type in column_items:
column_extras = []
if pk == column_name:
if column_name == single_pk:
column_extras.append("PRIMARY KEY")
if column_name in not_null:
column_extras.append("NOT NULL")
@ -233,12 +241,17 @@ class Database:
else "",
)
)
extra_pk = ""
if single_pk is None and pk and len(pk) > 1:
extra_pk = ",\n PRIMARY KEY ({pks})".format(
pks=", ".join(["[{}]".format(p) for p in pk])
)
columns_sql = ",\n".join(column_defs)
sql = """CREATE TABLE [{table}] (
{columns_sql}
){extra};
{columns_sql}{extra_pk}
);
""".format(
table=name, columns_sql=columns_sql, extra=extra
table=name, columns_sql=columns_sql, extra_pk=extra_pk
)
self.conn.execute(sql)
return self[name]
@ -382,6 +395,37 @@ class Table:
def pks(self):
return [column.name for column in self.columns if column.is_pk]
def get(self, pk_values):
if not isinstance(pk_values, (list, tuple)):
pk_values = [pk_values]
pks = self.pks
pk_names = []
if len(pks) == 0:
# rowid table
pk_names = ["rowid"]
last_pk = pk_values[0]
elif len(pks) == 1:
pk_names = [pks[0]]
last_pk = pk_values[0]
elif len(pks) > 1:
pk_names = pks
last_pk = pk_values
if len(pk_names) != len(pk_values):
raise NotFoundError(
"Need {} primary key value{}".format(
len(pk_names), "" if len(pk_names) == 1 else "s"
)
)
wheres = ["[{}] = ?".format(pk_name) for pk_name in pk_names]
rows = self.rows_where(" and ".join(wheres), pk_values)
try:
row = list(rows)[0]
self.last_pk = last_pk
return row
except IndexError:
raise NotFoundError
@property
def foreign_keys(self):
fks = []
@ -786,12 +830,13 @@ class Table:
self.last_pk = None
# self.last_rowid will be 0 if a "INSERT OR IGNORE" happened
if (hash_id or pk) and self.last_rowid:
self.last_pk = self.db.conn.execute(
"select [{}] from [{}] where rowid = ?".format(
hash_id or pk, self.name
),
(self.last_rowid,),
).fetchone()[0]
row = list(self.rows_where("rowid = ?", [self.last_rowid]))[0]
if hash_id:
self.last_pk = row[hash_id]
elif isinstance(pk, str):
self.last_pk = row[pk]
else:
self.last_pk = tuple(row[p] for p in pk)
return self
def upsert(