async def execute(self,
query,
params=None,
connection=None,
transaction=False):
if isinstance(query, (dict, list)):
# build zql query
query, params = build(query,
dialect=self.backend.dialect,
combine=True)
pool = await self.pool
connection = connection or self._connection
connection = aecho(connection) if connection else pool.acquire()
pquery = print_query(query, params)
async with connection as conn:
if self.prompt:
if not confirm(
f"{self.name} ({self.tag}): {SEP}{pquery}{SEPN}",
True):
raise Exception(f"{self}: execute aborted")
else:
self.log(f"{self}: execute{SEP}{pquery}{SEPN}")
transaction = conn.transaction() if transaction else aecho()
async with transaction:
try:
return await self.backend.execute(conn, query, params)
except Exception as e:
err = f"{self}: execute failed; {e.__class__.__name__}: {e}"
err += f"\nQuery:{SEP}{pquery}{SEPN}"
raise Exception(err)
async def stream(self,
query,
params=None,
transaction=True,
connection=None):
if isinstance(query, (dict, list)):
# build zql
queries = build(query, dialect=self.backend.dialect)
else:
queries = [(query, params)]
pool = await self.pool
connection = connection or self._connection
connection = aecho(connection) if connection else pool.acquire()
async with connection as conn:
transaction = conn.transaction() if transaction else aecho()
async with transaction:
for query, params in queries:
pquery = print_query(query, params)
if self.prompt:
if not confirm(
f"{self.name} ({self.tag}): {SEP}{pquery}{SEPN}",
True):
raise Exception(f"{self}: stream aborted")
else:
# self.log(f"{self}: stream{SEP}{pquery}{SEPN}")
async for row in self.backend.cursor(
conn, query, params):
yield row
async def query(self,
query,
params=None,
connection=None,
many=True,
columns=True,
transaction=False):
pool = await self.pool
connection = connection or self._connection
connection = aecho(connection) if connection else pool.acquire()
if isinstance(query, (dict, list)):
# build zql
queries = build(query, dialect=self.backend.dialect)
else:
queries = [(query, params)]
one = len(queries) == 1
all_results = []
async with connection as conn:
transaction = conn.transaction() if transaction else aecho()
async with transaction:
for query, params in queries:
pquery = print_query(query, params)
if self.prompt:
if not confirm(
f"{self.name} ({self.tag}): {SEP}{pquery}{SEPN}",
True):
raise Exception(f"{self}: query aborted")
else:
self.log(f"{self}: query{SEP}{pquery}{SEPN}")
try:
results = await self.backend.fetch(conn, query, params)
except Exception as e:
err = f"{self}: query failed; {e.__class__.__name__}: {e}"
err += f"\nQuery:{SEP}{pquery}{SEPN}"
raise Exception(err)
if many:
all_results.append(
results if columns else [r[0] for r in results])
continue
else:
num = len(results)
if num == 0:
# no results -> return None
all_results.append(None)
continue
if num != 1:
raise Exception(
f"{self}: query failed; expecting <=1 row, got {num}\n"
f"Query:{SEP}{query}{SEPN}")
result = results[0]
all_results.append(result if columns else result[0])
return all_results[0] if one else all_results
async def copy_from(self, **kwargs):
pool = await self.pool
table_name = kwargs.pop("table_name", None)
schema_name = kwargs.get('schema_name', None)
transaction = kwargs.pop("transaction", False)
connection = kwargs.pop("connection", self._connection)
connection = aecho(connection) if connection else pool.acquire()
close = kwargs.pop("close", False)
query = kwargs.pop("query", None)
params = kwargs.pop('params', None)
if table_name:
target_label = f"{schema_name}.{table_name}" if schema_name else table_name
else:
if not query:
raise NotImplementedError("table or query is required")
if isinstance(query, (list, dict)):
# compile from zQL
query, params = build(query,
dialect=self.backend.dialect,
combine=True)
target_label = print_query(query, params)
if self.prompt:
if not confirm(
f"{self.name} ({self.tag}): {SEP}copy from {target_label}{SEPN}",
True):
raise Exception(f"{self}: copy_from aborted")
else:
self.log(f"{self}: copy_from{SEP}{target_label}{SEPN}")
async with connection as conn:
transaction = conn.transaction() if transaction else aecho()
async with transaction:
result = None
if table_name:
result = await self.backend.copy_from_table(
conn, table_name, **kwargs)
else:
result = await self.backend.copy_from_query(
conn, query, params, **kwargs)
if close:
if hasattr(close, "close"):
# close passed in object
output = close
else:
# close output object
output = kwargs.get("output")
if getattr(output, "close"):
output.close()
return result
async def copy_to(self, **kwargs):
pool = await self.pool
table_name = kwargs.pop("table_name", None)
transaction = kwargs.pop("transaction", False)
schema_name = kwargs.get('schema_name', None)
columns = ', '.join(kwargs.get('columns', []))
columns = f' ({columns})' if columns else ''
target_label = f"{schema_name}.{table_name}{columns}" if schema_name else table_name
connection = kwargs.pop("connection", None) or self._connection
connection = aecho(connection) if connection else pool.acquire()
if self.prompt:
if not confirm(
f"{self.name} ({self.tag}): {SEP}copy to {target_label}{SEPN}",
True):
raise Exception(f"{self}: copy_to aborted")
else:
self.log(f"{self}: copy_to{SEP}{target_label}{SEPN}")
async with connection as conn:
transaction = conn.transaction() if transaction else aecho()
async with transaction:
return await self.backend.copy_to_table(
conn, table_name, **kwargs)
async def _copy_shard(
self,
source_model,
target_model,
target,
pk,
source_schema,
source_table,
target_schema,
target_table,
delete,
truncate,
cursor_min,
cursor_max,
):
# md5 check failed
# copy this shard
# maybe use transaction for delete?
# if there are no FK constraints, it shouldn't block
# but we would still get a more consistent table state
# for example, row counts will not change as large chunks of rows
# are deleted and re-added in a transaction
connection = None # await self.get_connection() if delete else None
transaction = aecho(
) # connection.transaction() if delete else aecho()
def get_query(q):
columns = q.table.order_by_alias(q.table.columns.keys())
q = q.take(*columns)
if cursor_min:
q = q.where({
"and": [{
'>': [pk, cursor_min]
}, {
'<=': [pk, cursor_max]
}]
})
else:
q = q.where({'<=': [pk, cursor_max]})
return q
try:
async with transaction:
source_query = get_query(source_model)
if delete:
if truncate:
await target_model.truncate()
else:
target_query = get_query(target_model)
await target_query.delete(connection=connection)
target_columns = target_model.table.order_by_alias(
target_model.table.columns.keys())
# copy from source to buffer
source_query = await source_query.get(zql=True)
if self.parallel_copy:
buffer = AsyncBuffer()
copy_from, copy_to = await gather(
self.copy_from(query=source_query,
output=buffer.write,
close=buffer),
target.copy_to(
table_name=target_table,
schema_name=target_schema,
source=buffer,
connection=connection,
columns=target_columns,
),
)
return copy_to
else:
buffer = io.BytesIO()
await self.copy_from(query=source_query, output=buffer)
buffer.seek(0)
return await target.copy_to(
table_name=target_table,
schema_name=target_schema,
source=buffer,
columns=target_columns,
connection=connection,
)
finally:
if connection:
await connection.close()