Modify pypgstac to pass mypy test

Author: lossyrob

pypgstac/pypgstac/pypgstac.py

@@ -1,8 +1,11 @@
 import asyncio
+from io import TextIOWrapper
 import os
-from typing import List
+import time
+from typing import Any, AsyncGenerator, Dict, Iterable, List, Optional, TypeVar, Union
 
 import asyncpg
+from asyncpg.connection import Connection
 import typer
 import orjson
 from smart_open import open
@@ -21,11 +24,11 @@
 migrations_dir = os.path.join(dirname, "migrations")
 
 
-def pglogger(conn, message):
+def pglogger(message: str) -> None:
     logging.debug(message)
 
 
-async def con_init(conn):
+async def con_init(conn: Connection) -> None:
     """Use orjson for json returns."""
     await conn.set_type_codec(
         "json",
@@ -42,33 +45,36 @@ async def con_init(conn):
 
 
 class DB:
-    pg_connection_string = None
-    connection = None
+    pg_connection_string: Optional[str] = None
+    connection: Optional[Connection] = None
 
-    def __init__(self, pg_connection_string: str = None):
+    def __init__(self, pg_connection_string: Optional[str] = None) -> None:
         self.pg_connection_string = pg_connection_string
 
-    async def create_connection(self):
-        self.connection = await asyncpg.connect(
+    async def create_connection(self) -> Connection:
+        connection: Connection = await asyncpg.connect(
             self.pg_connection_string,
             server_settings={
                 "search_path": "pgstac,public",
                 "application_name": "pypgstac",
             },
         )
-        await con_init(self.connection)
+        await con_init(connection)
+        self.connection = connection
         return self.connection
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> Connection:
         if self.connection is None:
             await self.create_connection()
+        assert self.connection is not None
         return self.connection
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.connection.close()
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        if self.connection:
+            await self.connection.close()
 
 
-async def run_migration(dsn: str = None):
+async def run_migration(dsn: Optional[str] = None) -> str:
     conn = await asyncpg.connect(dsn=dsn)
     async with conn.transaction():
         try:
@@ -102,27 +108,33 @@ async def run_migration(dsn: str = None):
             f"Pypgstac does not have a migration from {oldversion} to {version} ({migration_file})"
         )
 
-    with open(migration_file) as f:
-        migration_sql = f.read()
-        logging.debug(migration_sql)
-        async with conn.transaction():
-            conn.add_log_listener(pglogger)
-            await conn.execute(migration_sql)
-            await conn.execute(
-                """
-                INSERT INTO pgstac.migrations (version)
-                VALUES ($1);
-                """,
-                version,
-            )
-
-    await conn.close()
+    open_migration_file = open(migration_file)
+    if isinstance(open_migration_file, TextIOWrapper):
+        with open_migration_file as f:
+            migration_sql = f.read()
+            logging.debug(migration_sql)
+            async with conn.transaction():
+                conn.add_log_listener(pglogger)
+                await conn.execute(migration_sql)
+                await conn.execute(
+                    """
+                    INSERT INTO pgstac.migrations (version)
+                    VALUES ($1);
+                    """,
+                    version,
+                )
+
+        await conn.close()
+    else:
+        raise IOError(f"Unable to open {migration_file}")
     return version
 
 
 @app.command()
-def migrate(dsn: str = None):
-    typer.echo(asyncio.run(run_migration(dsn)))
+def migrate(dsn: Optional[str] = None) -> None:
+    """Migrate a pgstac database"""
+    version = asyncio.run(run_migration(dsn))
+    typer.echo(f'pgstac version {version}')
 
 
 class loadopt(str, Enum):
@@ -135,35 +147,42 @@ class tables(str, Enum):
     items = "items"
     collections = "collections"
 
+# Types of iterable that load_iterator can support
+T = TypeVar('T', Iterable[bytes], Iterable[Dict[str, Any]], Iterable[str])
+
 
-async def aiter(list: List):
-    for i in list:
-        if isinstance(i, bytes):
-            i = i.decode("utf-8")
-        elif isinstance(i, dict):
-            i = orjson.dumps(i).decode("utf-8")
-        if isinstance(i, str):
-            line = "\n".join(
-                [
-                    i.rstrip()
-                    .replace(r"\n", r"\\n")
-                    .replace(r"\t", r"\\t")
-                ]
-            ).encode("utf-8")
-            yield line
+async def aiter(list: T) -> AsyncGenerator[bytes, None]:
+    for item in list:
+        item_str: str
+        if isinstance(item, bytes):
+            item_str = item.decode("utf-8")
+        elif isinstance(item, dict):
+            item_str = orjson.dumps(item).decode("utf-8")
+        elif isinstance(item, str):
+            item_str = item
         else:
-            raise Exception(f"Could not parse {i}")
+            raise ValueError(f"Cannot load iterator with values of type {type(item)} (value {item})")
+
 
+        line = "\n".join(
+            [
+                item_str.rstrip()
+                .replace(r"\n", r"\\n")
+                .replace(r"\t", r"\\t")
+            ]
+        ).encode("utf-8")
+        yield line
 
-async def copy(iter, table: tables, conn: asyncpg.Connection):
+
+async def copy(iter: T, table: tables, conn: asyncpg.Connection) -> None:
     logger.debug(f"copying to {table} directly")
     logger.debug(f"iter: {iter}")
-    iter = aiter(iter)
+    bytes_iter = aiter(iter)
     async with conn.transaction():
         logger.debug("Copying data")
         await conn.copy_to_table(
             table,
-            source=iter,
+            source=bytes_iter,
             columns=["content"],
             format="csv",
             quote=chr(27),
@@ -179,10 +198,10 @@ async def copy(iter, table: tables, conn: asyncpg.Connection):
 
 
 async def copy_ignore_duplicates(
-    iter, table: tables, conn: asyncpg.Connection
-):
+    iter: T, table: tables, conn: asyncpg.Connection
+) -> None:
     logger.debug(f"inserting to {table} ignoring duplicates")
-    iter = aiter(iter)
+    bytes_iter = aiter(iter)
     async with conn.transaction():
         await conn.execute(
             """
@@ -192,7 +211,7 @@ async def copy_ignore_duplicates(
         )
         await conn.copy_to_table(
             "pgstactemp",
-            source=iter,
+            source=bytes_iter,
             columns=["content"],
             format="csv",
             quote=chr(27),
@@ -218,9 +237,9 @@ async def copy_ignore_duplicates(
         logger.debug("Data Inserted")
 
 
-async def copy_upsert(iter, table: tables, conn: asyncpg.Connection):
+async def copy_upsert(iter: T, table: tables, conn: asyncpg.Connection) -> None:
     logger.debug(f"upserting to {table}")
-    iter = aiter(iter)
+    bytes_iter = aiter(iter)
     async with conn.transaction():
         await conn.execute(
             """
@@ -230,7 +249,7 @@ async def copy_upsert(iter, table: tables, conn: asyncpg.Connection):
         )
         await conn.copy_to_table(
             "pgstactemp",
-            source=iter,
+            source=bytes_iter,
             columns=["content"],
             format="csv",
             quote=chr(27),
@@ -258,24 +277,28 @@ async def copy_upsert(iter, table: tables, conn: asyncpg.Connection):
 
 
 async def load_iterator(
-    iter, table: tables, conn: asyncpg.Connection, method: loadopt = "insert"
+    iter: T, table: tables, conn: asyncpg.Connection, method: loadopt = loadopt.insert
 ):
     logger.debug(f"Load Iterator Connection: {conn}")
-    if method == "insert":
+    if method == loadopt.insert:
         await copy(iter, table, conn)
-    elif method == "insert_ignore":
+    elif method == loadopt.insert_ignore:
         await copy_ignore_duplicates(iter, table, conn)
     else:
         await copy_upsert(iter, table, conn)
 
 
 async def load_ndjson(
-    file: str, table: tables, method: loadopt = "insert", dsn: str = None
-):
+    file: str, table: tables, method: loadopt = loadopt.insert, dsn: str = None
+) -> None:
     print(f"loading {file} into {table} using {method}")
-    with open(file, "rb") as f:
-        async with DB(dsn) as conn:
-            await load_iterator(f, table, conn, method)
+    open_file = open(file, "rb")
+    if isinstance(open_file, TextIOWrapper):
+        with open_file as f:
+            async with DB(dsn) as conn:
+                await load_iterator(f, table, conn, method)
+    else:
+        raise IOError(f"Cannot read {file}")
 
 
 @app.command()
@@ -286,13 +309,38 @@ def load(
     method: loadopt = typer.Option(
         "insert", prompt="How to deal conflicting ids"
     ),
-):
+) -> None:
+    "Load STAC data into a pgstac database."
     typer.echo(
         asyncio.run(
             load_ndjson(file=file, table=table, dsn=dsn, method=method)
         )
     )
 
+@app.command()
+def pgready(dsn: Optional[str] = None) -> None:
+    """Wait for a pgstac database to accept connections"""
+    async def wait_on_connection() -> bool:
+        cnt = 0
+
+        print("Waiting for pgstac to come online...", end="", flush=True)
+        while True:
+            if cnt > 150:
+                raise Exception("Unable to connect to database")
+            try:
+                print(".", end="", flush=True)
+                conn = await asyncpg.connect()
+                await conn.execute("SELECT 1")
+                await conn.close()
+                print("success!")
+                return True
+            except Exception:
+                time.sleep(0.1)
+                cnt += 1
+
+
+    asyncio.run(wait_on_connection())
+
 
 if __name__ == "__main__":
     app()