Duckdb #75
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,5 +1,5 @@ | ||
| FROM python:3.13.3-alpine3.22 | ||
| RUN apk add bash poetry build-base | ||
| WORKDIR /app | ||
| ADD . /app | ||
| RUN mkdir /pypoetry | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,38 +1,186 @@ | ||
| """Data dumping functions.""" | ||
| import csv | ||
| import io | ||
| from abc import ABC, abstractmethod | ||
| from pathlib import Path | ||
|
|
||
| import pandas as pd | ||
| import sqlalchemy | ||
| from sqlalchemy.schema import MetaData | ||
|
|
||
| from datafaker.utils import create_db_engine, get_sync_engine, logger | ||
|
|
||
|
|
||
| class TableWriter(ABC): | ||
| """Writes a table out to a file.""" | ||
|
|
||
| EXTENSION = ".csv" | ||
|
|
||
| def __init__(self, metadata: MetaData, dsn: str, schema: str | None) -> None: | ||
| """ | ||
| Initialize the TableWriter. | ||
|
|
||
| :param metadata: The metadata for our database. | ||
| :param dsn: The connection string for our database. | ||
| :param schema: The schema name for our database, or None for the default. | ||
| """ | ||
| self._metadata = metadata | ||
| self._dsn = dsn | ||
| self._schema = schema | ||
|
|
||
| def connect(self) -> sqlalchemy.engine.Connection: | ||
| """Connect to the database.""" | ||
| engine = get_sync_engine(create_db_engine(self._dsn, schema_name=self._schema)) | ||
| return engine.connect() | ||
|
|
||
| @abstractmethod | ||
| def write_file(self, table: sqlalchemy.Table, filepath: Path) -> bool: | ||
| """ | ||
| Write the named table into the named file. | ||
|
|
||
| :param table: The table to output | ||
| :param dir: The directory to write into. | ||
| :return: ``true`` on success, otherwise ``false``. | ||
| """ | ||
|
|
||
| def write(self, table: sqlalchemy.Table, directory: Path) -> bool: | ||
| """ | ||
| Write the table into a directory with a filename based on the table's name. | ||
|
|
||
| :param table: The table to write out. | ||
| :param directory: The directory to write the table into. | ||
| :return: ``true`` on success, otherwise ``false``. | ||
| """ | ||
| tn = table.name | ||
| # DuckDB tables derived from files have confusing suffixes | ||
| # that we should probably remove | ||
| tn = tn.removesuffix(".csv") | ||
| tn = tn.removesuffix(".parquet") | ||
| return self.write_file(table, directory / f"{tn}{self.EXTENSION}") | ||
|
|
||
|
|
||
| class ParquetTableWriter(TableWriter): | ||
| """Writes the table to a Parquet file.""" | ||
|
|
||
| EXTENSION = ".parquet" | ||
|
|
||
| def write_file(self, table: sqlalchemy.Table, filepath: Path) -> bool: | ||
| """ | ||
| Write the named table into the named file. | ||
|
|
||
| :param table: The table to output | ||
| :param filename: The filename of the file to write to. | ||
| :return: ``true`` on success, otherwise ``false``. | ||
| """ | ||
| with self.connect() as connection: | ||
| dates = [ | ||
| str(name) | ||
| for name, col in table.columns.items() | ||
| if isinstance( | ||
| col.type, | ||
| ( | ||
| sqlalchemy.types.DATE, | ||
| sqlalchemy.types.DATETIME, | ||
| sqlalchemy.types.TIMESTAMP, | ||
| ), | ||
| ) | ||
| ] | ||
| df = pd.read_sql( | ||
| sql=f"SELECT * FROM {table.name}", | ||
| con=connection, | ||
| columns=[str(col.name) for col in table.columns.values()], | ||
| parse_dates=dates, | ||
| ) | ||
| df.to_parquet(filepath) | ||
| return True | ||
|
|
||
|
|
||
| class DuckDbParquetTableWriter(ParquetTableWriter): | ||
| """ | ||
| Writes the table to a Parquet file using DuckDB SQL. | ||
|
|
||
| The Pandas method used by ParquetTableWriter currently | ||
| does not work with DuckDB. | ||
| """ | ||
|
|
||
| def write_file(self, table: sqlalchemy.Table, filepath: Path) -> bool: | ||
| """ | ||
| Write the named table into the named file. | ||
|
|
||
| :param table: The table to output | ||
| :param filename: The filename of the file to write to. | ||
| :return: ``true`` on success, otherwise ``false``. | ||
| """ | ||
| with self.connect() as connection: | ||
| result = connection.execute( | ||
| sqlalchemy.text( | ||
| # We need the double quotes to get DuckDB to read the table not the file. | ||
| f"COPY \"{table.name}\" TO '{filepath}' (FORMAT PARQUET)" | ||
| ) | ||
| ) | ||
| return result is not None | ||
|
|
||
|
|
||
| def get_parquet_table_writer( | ||
| metadata: MetaData, dsn: str, schema: str | None | ||
| ) -> TableWriter: | ||
| """ | ||
| Get a ``TableWriter`` that writes parquet files. | ||
|
|
||
| :param metadata: The database metadata containing the tables to be dumped to files. | ||
| :param dsn: The database connection string. | ||
| :param schema: The schema name, if required. | ||
| :return: ``TableWriter`` to write a parquet file. | ||
| """ | ||
| if dsn.startswith("duckdb:"): | ||
| return DuckDbParquetTableWriter(metadata, dsn, schema) | ||
| return ParquetTableWriter(metadata, dsn, schema) | ||
|
|
||
|
|
||
| class TableWriterIO(TableWriter): | ||
| """Writes the table to an output object.""" | ||
|
|
||
| @abstractmethod | ||
| def write_io(self, table: sqlalchemy.Table, out: io.TextIOBase) -> bool: | ||
| """ | ||
| Write the named table into the named file. | ||
|
|
||
| :param table: The table to output | ||
| :param filename: The filename of the file to write to. | ||
| :return: ``true`` on success, otherwise ``false``. | ||
| """ | ||
|
|
||
| def write_file(self, table: sqlalchemy.Table, filepath: Path) -> bool: | ||
| """ | ||
| Write the named table into the named file. | ||
|
|
||
| :param table: The table to output | ||
| :param filename: The filename of the file to write to. | ||
| :return: ``true`` on success, otherwise ``false``. | ||
| """ | ||
| with open(filepath, "wt", newline="", encoding="utf-8") as out: | ||
| return self.write_io(table, out) | ||
|
|
||
|
|
||
| class CsvTableWriter(TableWriterIO): | ||
| """Writes the table to a CSV file.""" | ||
|
|
||
| def write_io(self, table: sqlalchemy.Table, out: io.TextIOBase) -> bool: | ||
| """ | ||
| Write the named table into the named file. | ||
|
|
||
| :param table: The table to output | ||
| :param filename: The filename of the file to write to. | ||
| :return: ``True`` on success, otherwise ``False``. | ||
| """ | ||
| if table.name not in self._metadata.tables: | ||
| logger.error("%s is not a table described in the ORM file", table.name) | ||
| return False | ||
| table = self._metadata.tables[table.name] | ||
| csv_out = csv.writer(out, quoting=csv.QUOTE_MINIMAL) | ||
| csv_out.writerow(table.columns.keys()) | ||
| with self.connect() as connection: | ||
| result = connection.execute(sqlalchemy.select(table)) | ||
| for row in result: | ||
| csv_out.writerow(row) | ||
| return True |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ooh this is fun
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pretty nasty actually. But yes, fun that this hook exists!