diff --git a/environment.yml b/environment.yml index 4947ac6..ba72ed3 100644 --- a/environment.yml +++ b/environment.yml @@ -38,6 +38,7 @@ dependencies: - setuptools==70.3.0 - pyparsing==3.1.4 - findlibs==0.1.2 + - psycopg2-binary==2.9.9 - . - ./weather_dl - ./weather_mv diff --git a/setup.py b/setup.py index cdbf585..c799f4b 100644 --- a/setup.py +++ b/setup.py @@ -41,6 +41,7 @@ "google-cloud-firestore", "firebase-admin", "urllib3==1.26.5", + "psycopg2-binary", ] weather_mv_requirements = [ diff --git a/weather_dl/README.md b/weather_dl/README.md index 96b8310..08575d2 100644 --- a/weather_dl/README.md +++ b/weather_dl/README.md @@ -46,8 +46,9 @@ _Common options_: (stdout for direct runner). You can set the name of the manifest as the hostname of a URL with the 'cli' protocol. For example, `cli://manifest` will prefix all the manifest logs as '[manifest]'. In addition, users can specify either a Firestore collection URI (`fs://?projectId=`), or - BigQuery table (`bq://..`), or `noop://` - for an in-memory location. + BigQuery table (`bq://..`), or + PostgreSQL table (`postgresql://:@:/?table=`), or + `noop://` for an in-memory location. * `-n, --num-requests-per-key`: Number of concurrent requests to make per API key. Default: make an educated guess per client & config. Please see the client documentation for more details. * `-p, --partition-chunks`: Group shards into chunks of this size when computing the partitions. Specifically, this @@ -63,6 +64,10 @@ _Common options_: > Note: > * In case of BigQuery manifest tool will create the BQ table itself, if not already present. > Or it will use the existing table but can report errors in case of schema mismatch. +> * In case of PostgreSQL manifest tool will create the table (and an index on `config_name`) itself, if not +> already present. The table is upserted on the `location` column, so the latest row per location represents +> the current state of a download. This works with any PostgreSQL-compatible database, including +> [Cloud Spanner's PostgreSQL interface](https://cloud.google.com/spanner/docs/postgresql-interface). > * To run complex queries on the Firestore manifest, users may find it helpful to replicate Firestore to BigQuery > using the automated process described in > [this article](https://medium.com/@ammppp/automated-firestore-replication-to-bigquery-15915d518e38). diff --git a/weather_dl/download_pipeline/manifest.py b/weather_dl/download_pipeline/manifest.py index d2a7255..8aef4d5 100644 --- a/weather_dl/download_pipeline/manifest.py +++ b/weather_dl/download_pipeline/manifest.py @@ -22,11 +22,12 @@ import logging import os import pandas as pd +import re import threading import time import traceback import typing as t -from urllib.parse import urlparse, parse_qsl +from urllib.parse import urlparse, parse_qsl, urlencode, urlunparse from .util import ( to_json_serializable_type, @@ -577,6 +578,183 @@ def _update(self, download_status: DownloadStatus) -> None: logger.debug(download_status) +class PostgresManifest(Manifest): + """A PostgreSQL Manifest. + + This implementation stores download statuses in a relational table, where + each row is keyed by the download `location`. Updating a status upserts the + corresponding row, so the latest value represents the current state of a + download. + + Location format: + postgresql://:@:/?table= + + The `table` query parameter is optional and defaults to 'manifest'. The + table and an index on `config_name` are created if they do not already + exist. Any additional query parameters are forwarded to `psycopg2` as + connection options (e.g. `sslmode=require`). + + Example: + >>> manifest = PostgresManifest( + ... Location('postgresql://user:pass@localhost:5432/weather?table=downloads')) + """ + + DEFAULT_TABLE = 'manifest' + + _COLUMNS = ( + 'config_name', 'dataset', 'selection', 'location', 'area', 'stage', + 'status', 'error', 'username', 'size', 'scheduled_time', + 'retrieve_start_time', 'retrieve_end_time', 'fetch_start_time', + 'fetch_end_time', 'download_start_time', 'download_end_time', + 'upload_start_time', 'upload_end_time', + ) + + _TIMESTAMP_COLUMNS = ( + 'scheduled_time', 'retrieve_start_time', 'retrieve_end_time', + 'fetch_start_time', 'fetch_end_time', 'download_start_time', + 'download_end_time', 'upload_start_time', 'upload_end_time', + ) + + _CREATE_TABLE_TEMPLATE = """ + CREATE TABLE IF NOT EXISTS {table} ( + config_name TEXT NOT NULL, + dataset TEXT, + selection TEXT NOT NULL, + location TEXT PRIMARY KEY, + area TEXT, + stage TEXT, + status TEXT NOT NULL, + error TEXT, + username TEXT NOT NULL, + size DOUBLE PRECISION, + scheduled_time TIMESTAMPTZ, + retrieve_start_time TIMESTAMPTZ, + retrieve_end_time TIMESTAMPTZ, + fetch_start_time TIMESTAMPTZ, + fetch_end_time TIMESTAMPTZ, + download_start_time TIMESTAMPTZ, + download_end_time TIMESTAMPTZ, + upload_start_time TIMESTAMPTZ, + upload_end_time TIMESTAMPTZ + ) + """ + + def __init__(self, location: Location) -> None: + super().__init__(location) + self.dsn, self.table = self.parse_location(location) + self._ensure_table() + + @classmethod + def parse_location(cls, location: str) -> t.Tuple[str, str]: + """Split a manifest location into a PostgreSQL DSN and table name. + + The `table` query parameter is removed from the DSN, since it is not a + valid libpq connection option. + + Raises: + ManifestException: If the table name is not a valid SQL identifier. + """ + parsed = urlparse(location) + query_params = dict(parse_qsl(parsed.query, keep_blank_values=True)) + table = query_params.pop('table', cls.DEFAULT_TABLE) + + if not re.fullmatch(r'[A-Za-z_][A-Za-z0-9_]{0,62}', table): + raise ManifestException( + f"Invalid table name {table!r}. Table names must be valid PostgreSQL identifiers.") + + dsn = urlunparse(parsed._replace(query=urlencode(query_params))) + return dsn, table + + def _connect(self): + """Open a new connection to the manifest's PostgreSQL database.""" + import psycopg2 + return psycopg2.connect(self.dsn) + + def _ensure_table(self) -> None: + """Create the manifest table and its index if they don't exist.""" + import psycopg2.sql as sql + create_table = sql.SQL(self._CREATE_TABLE_TEMPLATE).format(table=sql.Identifier(self.table)) + create_index = sql.SQL( + 'CREATE INDEX IF NOT EXISTS {index} ON {table} (config_name)' + ).format( + index=sql.Identifier(f'{self.table}_config_name_idx'), + table=sql.Identifier(self.table), + ) + + conn = self._connect() + try: + with conn, conn.cursor() as cursor: + cursor.execute(create_table) + cursor.execute(create_index) + finally: + conn.close() + + def _read(self, location: str) -> DownloadStatus: + """Read the download status for a location from the manifest.""" + import psycopg2.sql as sql + query = sql.SQL('SELECT {columns} FROM {table} WHERE location = %s').format( + columns=sql.SQL(', ').join(map(sql.Identifier, self._COLUMNS)), + table=sql.Identifier(self.table), + ) + + conn = self._connect() + try: + with conn, conn.cursor() as cursor: + cursor.execute(query, (location,)) + row = cursor.fetchone() + if row is None: + return DownloadStatus() + record = dict(zip(self._COLUMNS, row)) + finally: + conn.close() + + return DownloadStatus.from_dict(self._to_serializable(record)) + + def _update(self, download_status: DownloadStatus) -> None: + """Upsert the download status for a location into the manifest.""" + import psycopg2.sql as sql + status = DownloadStatus.to_dict(download_status) + values = [self._to_db_value(column, status.get(column)) for column in self._COLUMNS] + + query = sql.SQL( + 'INSERT INTO {table} ({columns}) VALUES ({values}) ' + 'ON CONFLICT (location) DO UPDATE SET {updates}' + ).format( + table=sql.Identifier(self.table), + columns=sql.SQL(', ').join(map(sql.Identifier, self._COLUMNS)), + values=sql.SQL(', ').join([sql.Placeholder()] * len(self._COLUMNS)), + updates=sql.SQL(', ').join( + sql.SQL('{column} = EXCLUDED.{column}').format(column=sql.Identifier(column)) + for column in self._COLUMNS if column != 'location' + ), + ) + + conn = self._connect() + try: + with conn, conn.cursor() as cursor: + cursor.execute(query, values) + finally: + conn.close() + + logger.debug('Manifest written to.') + logger.debug(download_status) + + @classmethod + def _to_db_value(cls, column: str, value: t.Any) -> t.Any: + """Normalize a DownloadStatus field before writing it to PostgreSQL.""" + if column in cls._TIMESTAMP_COLUMNS and not value: + return None + return value + + @classmethod + def _to_serializable(cls, record: t.Dict) -> t.Dict: + """Convert a PostgreSQL row into a manifest-ingestible dictionary.""" + return { + key: value.isoformat() if isinstance(value, datetime.datetime) else value + for key, value in record.items() + } + + class FirestoreManifest(Manifest): """A Firestore Manifest. This Manifest implementation stores DownloadStatuses in a Firebase document store. @@ -720,6 +898,8 @@ def _update(self, download_status: DownloadStatus) -> None: 'cli': ConsoleManifest, 'fs': FirestoreManifest, 'bq': BQManifest, + 'postgresql': PostgresManifest, + 'postgres': PostgresManifest, '': LocalManifest, }) diff --git a/weather_dl/download_pipeline/manifest_test.py b/weather_dl/download_pipeline/manifest_test.py index a36e0be..fd06492 100644 --- a/weather_dl/download_pipeline/manifest_test.py +++ b/weather_dl/download_pipeline/manifest_test.py @@ -13,13 +13,30 @@ # limitations under the License. import json +import os import random import string import tempfile import typing as t import unittest -from .manifest import LocalManifest, Location, DownloadStatus, Status, Stage +from .manifest import ( + LocalManifest, + Location, + DownloadStatus, + ManifestException, + PostgresManifest, + Status, + Stage, +) + +POSTGRES_DSN = os.environ.get('POSTGRES_MANIFEST_DSN') +POSTGRES_TEST_TABLE = 'weather_tools_manifest_test' + + +def postgres_location(table: str) -> Location: + separator = '&' if '?' in (POSTGRES_DSN or '') else '?' + return Location(f'{POSTGRES_DSN}{separator}table={table}') def rand_str(max_len=32): @@ -99,3 +116,125 @@ def is_valid_json(self, file: t.IO) -> None: json.dumps(json.load(file)) except json.JSONDecodeError: self.fail('JSON is invalid.') + + +class PostgresManifestLocationTest(unittest.TestCase): + + def test_parse_location__extracts_table(self): + dsn, table = PostgresManifest.parse_location( + Location('postgresql://user:pass@localhost:5432/weather?table=downloads')) + + self.assertEqual(dsn, 'postgresql://user:pass@localhost:5432/weather') + self.assertEqual(table, 'downloads') + + def test_parse_location__defaults_table(self): + dsn, table = PostgresManifest.parse_location(Location('postgresql://user:pass@localhost/weather')) + + self.assertEqual(dsn, 'postgresql://user:pass@localhost/weather') + self.assertEqual(table, PostgresManifest.DEFAULT_TABLE) + + def test_parse_location__preserves_other_query_params(self): + dsn, table = PostgresManifest.parse_location( + Location('postgresql://user:pass@localhost/weather?table=downloads&sslmode=require')) + + self.assertEqual(table, 'downloads') + self.assertIn('sslmode=require', dsn) + self.assertNotIn('table=', dsn) + + def test_parse_location__rejects_invalid_table_name(self): + with self.assertRaises(ManifestException): + PostgresManifest.parse_location(Location('postgresql://localhost/weather?table=bad;drop')) + + def test_manifest_registry__supports_postgres_schemes(self): + from .manifest import MANIFESTS + + self.assertIs(MANIFESTS['postgresql'], PostgresManifest) + self.assertIs(MANIFESTS['postgres'], PostgresManifest) + + def test_to_db_value__normalizes_empty_timestamps(self): + self.assertIsNone(PostgresManifest._to_db_value('scheduled_time', '')) + self.assertIsNone(PostgresManifest._to_db_value('upload_end_time', None)) + self.assertEqual( + PostgresManifest._to_db_value('scheduled_time', '2023-02-07T17:15:26+00:00'), + '2023-02-07T17:15:26+00:00') + + def test_to_db_value__leaves_non_timestamps_untouched(self): + self.assertEqual(PostgresManifest._to_db_value('username', ''), '') + + def test_to_serializable__converts_datetimes_to_isoformat(self): + import datetime + + record = { + 'location': 'gs://bucket/file.nc', + 'scheduled_time': datetime.datetime(2023, 2, 7, 17, 15, 26, tzinfo=datetime.timezone.utc), + 'size': 0.5, + } + + actual = PostgresManifest._to_serializable(record) + + self.assertEqual(actual['scheduled_time'], '2023-02-07T17:15:26+00:00') + self.assertEqual(actual['location'], 'gs://bucket/file.nc') + self.assertEqual(actual['size'], 0.5) + + +@unittest.skipUnless(POSTGRES_DSN, 'POSTGRES_MANIFEST_DSN is not set; skipping PostgreSQL tests.') +class PostgresManifestTest(unittest.TestCase): + + def setUp(self): + self._drop_table() + self.manifest = PostgresManifest(postgres_location(POSTGRES_TEST_TABLE)) + + def tearDown(self): + self._drop_table() + + def _drop_table(self): + import psycopg2 + dsn, table = PostgresManifest.parse_location(postgres_location(POSTGRES_TEST_TABLE)) + with psycopg2.connect(dsn) as conn, conn.cursor() as cursor: + cursor.execute(f'DROP TABLE IF EXISTS "{table}"') + + def test_read__missing_location_returns_empty_status(self): + actual = self.manifest._read('gs://bucket/does-not-exist.nc') + + self.assertEqual(actual, DownloadStatus()) + + def test_update__round_trips_download_status(self): + expected = make_download_status(location='gs://bucket/file.nc') + self.manifest._update(expected) + + actual = self.manifest._read('gs://bucket/file.nc') + + self.assertEqual(actual.location, expected.location) + self.assertEqual(actual.status, expected.status) + self.assertEqual(actual.stage, expected.stage) + self.assertEqual(actual.username, expected.username) + self.assertEqual(actual.size, expected.size) + self.assertEqual(actual.scheduled_time, expected.scheduled_time) + self.assertEqual(actual.download_start_time, expected.download_start_time) + self.assertEqual(actual.upload_end_time, expected.upload_end_time) + + def test_update__overwrites_existing_status(self): + location = 'gs://bucket/file.nc' + self.manifest._update(make_download_status(location=location)) + expected = DownloadStatus(location=location, status=Status.FAILURE, error='boom', username='alice') + self.manifest._update(expected) + + actual = self.manifest._read(location) + + self.assertEqual(actual.location, location) + self.assertEqual(actual.status, Status.FAILURE) + self.assertEqual(actual.error, 'boom') + self.assertEqual(actual.username, 'alice') + + def test_update__keeps_one_row_per_location(self): + import psycopg2 + location = 'gs://bucket/file.nc' + for _ in range(5): + self.manifest._update(make_download_status(location=location)) + + dsn, table = PostgresManifest.parse_location(postgres_location(POSTGRES_TEST_TABLE)) + with psycopg2.connect(dsn) as conn, conn.cursor() as cursor: + cursor.execute(f'SELECT COUNT(*) FROM "{table}" WHERE location = %s', (location,)) + count = cursor.fetchone()[0] + + self.assertEqual(count, 1) diff --git a/weather_dl/download_pipeline/pipeline.py b/weather_dl/download_pipeline/pipeline.py index 3ce74f0..79f021a 100644 --- a/weather_dl/download_pipeline/pipeline.py +++ b/weather_dl/download_pipeline/pipeline.py @@ -155,7 +155,10 @@ def run(argv: t.List[str], save_main_session: bool = True) -> PipelineArgs: "('fs://?projectId='), or BigQuery table " "('bq://..') [Note: Tool will create the BQ table " "itself, if not already present. Or it will use the existing table but can report errors " - "in case of schema mismatch.], or 'noop://' for an in-memory location.") + "in case of schema mismatch.], or PostgreSQL table " + "('postgresql://:@:/?table=') " + "[Note: Tool will create the PostgreSQL table itself, if not already present.], or " + "'noop://' for an in-memory location.") parser.add_argument('-n', '--num-requests-per-key', type=int, default=-1, help='Number of concurrent requests to make per API key. ' 'Default: make an educated guess per client & config. ' diff --git a/weather_dl/setup.py b/weather_dl/setup.py index de6777c..5f02623 100644 --- a/weather_dl/setup.py +++ b/weather_dl/setup.py @@ -40,6 +40,7 @@ "urllib3==1.26.5", "google-cloud-firestore==2.6.0", "firebase-admin==6.0.1", + "psycopg2-binary", # "gcloud" should already be installed in the host image. # If we install it here, we'll hit auth issues. ]