diff --git a/apps/cli/utilities/core_sampler/core_sampler/core_sampler.py b/apps/cli/utilities/core_sampler/core_sampler/core_sampler.py index f35137ed2daa288948a86f1437a7a0ac0bbaadca..9ee27d54ae85bbe3bb2c0bed2d742ac52416252a 100644 --- a/apps/cli/utilities/core_sampler/core_sampler/core_sampler.py +++ b/apps/cli/utilities/core_sampler/core_sampler/core_sampler.py @@ -9,6 +9,7 @@ The core sampler outputs an SQL file you can use to load the core sample into a import argparse from typing import Any, Iterable +import datetime import psycopg2 as pg import psycopg2.extras as extras @@ -54,7 +55,7 @@ class MDDBConnector: class RowWriter: - def write_row(self, table: str, row: dict): + def write_rows(self, table: str, row: dict): raise NotImplementedError @@ -106,8 +107,24 @@ class CoreSampler(RowWriter): """ rows.write_to(self) - def write_row(self, table: str, row: dict): - pass # print((table, row)) + def write_rows(self, table: str, rows: list[dict]): + columns = rows[0].keys() + print(f"COPY {table} ({', '.join(columns)}) FROM stdin;") + for row in rows: + print("\t".join([self.copy_format(row[col]) for col in columns])) + print("\.") + + def copy_format(self, value): + if value == None: + return "\\N" + elif isinstance(value, str): + return value + elif isinstance(value, int) or isinstance(value, float): + return str(value) + elif isinstance(value, datetime.date): + return value.isoformat() + else: + raise TypeError(f"Unable to figure out what to do with {value} of type {type(value)}") class Table: @@ -163,7 +180,7 @@ class Table: :param columns: columns to consider in the generated WHERE clause :return: RowSet for the rows in this table """ - + # print(f"Fetching from {self.name} according to {','.join(columns)}") # 1. Escape the WHERE clause entries. it's important to use the primary keys # to retrieve the values from the previous resultset; the new columns will # appear in the query below. @@ -243,8 +260,7 @@ class RowSet: return self.table.relations() def write_to(self, writer: RowWriter): - for row in self.rows: - writer.write_row(self.table.name, row) + writer.write_rows(self.table.name, self.rows) def __iter__(self): return iter(self.rows)