chore: fix formatting in extension tests

projects/extension/tests/conftest.py

@@ -29,7 +29,9 @@ def does_test_db_exist(cur: psycopg.Cursor) -> bool:
 
 
 def drop_test_db(cur: psycopg.Cursor) -> None:
-    cur.execute("select pg_terminate_backend(pid) from pg_stat_activity where datname = 'test'")
+    cur.execute(
+        "select pg_terminate_backend(pid) from pg_stat_activity where datname = 'test'"
+    )
     cur.execute("drop database test with (force)")
 
 
@@ -42,12 +44,16 @@ def create_test_db(cur: psycopg.Cursor) -> None:
 @pytest.fixture(scope="session", autouse=True)
 def set_up_test_db() -> None:
     # create a test user and test database owned by the test user
-    with psycopg.connect("postgres://[email protected]:5432/postgres", autocommit=True) as con:
+    with psycopg.connect(
+        "postgres://[email protected]:5432/postgres", autocommit=True
+    ) as con:
         with con.cursor() as cur:
             create_test_user(cur)
             create_test_db(cur)
     # grant some things to the test user in the test database
-    with psycopg.connect("postgres://[email protected]:5432/test", autocommit=True) as con:
+    with psycopg.connect(
+        "postgres://[email protected]:5432/test", autocommit=True
+    ) as con:
         with con.cursor() as cur:
             cur.execute("grant execute on function pg_read_binary_file(text) to test")
             cur.execute("grant pg_read_server_files to test")

projects/extension/tests/contents/test_contents.py

@@ -30,14 +30,16 @@ def host_dir() -> Path:
 
 
 def init() -> None:
-    cmd = " ".join([
-        "psql",
-        f'''-d "{db_url("postgres", "postgres")}"''',
-        "-v ON_ERROR_STOP=1",
-        "-X",
-        f"-o {docker_dir()}/output.actual",
-        f"-f {docker_dir()}/init.sql",
-    ])
+    cmd = " ".join(
+        [
+            "psql",
+            f'''-d "{db_url("postgres", "postgres")}"''',
+            "-v ON_ERROR_STOP=1",
+            "-X",
+            f"-o {docker_dir()}/output.actual",
+            f"-f {docker_dir()}/init.sql",
+        ]
+    )
     if where_am_i() != "docker":
         cmd = f"docker exec -w {docker_dir()} pgai {cmd}"
     subprocess.run(cmd, check=True, shell=True, env=os.environ, cwd=str(host_dir()))
@@ -47,4 +49,4 @@ def test_contents() -> None:
     init()
     actual = host_dir().joinpath("output.actual").read_text()
     expected = host_dir().joinpath("output.expected").read_text()
-    assert actual == expected
+    assert actual == expected

projects/extension/tests/dump_restore/test_dump_restore.py

@@ -33,32 +33,41 @@ def host_dir() -> Path:
 
 
 def create_user() -> None:
-    with psycopg.connect(db_url(user="postgres", dbname="postgres"), autocommit=True) as con:
+    with psycopg.connect(
+        db_url(user="postgres", dbname="postgres"), autocommit=True
+    ) as con:
         with con.cursor() as cur:
-            cur.execute("""
+            cur.execute(
+                """
                 select count(*) > 0
                 from pg_catalog.pg_roles
                 where rolname = %s
-            """, (USER,))
+            """,
+                (USER,),
+            )
             exists: bool = cur.fetchone()[0]
             if not exists:
                 cur.execute(f"create user {USER}")  # NOT a superuser
 
 
 def create_database(dbname: str) -> None:
-    with psycopg.connect(db_url(user="postgres", dbname="postgres"), autocommit=True) as con:
+    with psycopg.connect(
+        db_url(user="postgres", dbname="postgres"), autocommit=True
+    ) as con:
         with con.cursor() as cur:
             cur.execute(f"drop database if exists {dbname} with (force)")
             cur.execute(f"create database {dbname} with owner {USER}")
 
 
 def dump_db() -> None:
     host_dir().joinpath("dump.sql").unlink(missing_ok=True)
-    cmd = " ".join([
-        "pg_dump -Fp",
-        f'''-d "{db_url(USER, "src")}"''',
-        f'''-f {docker_dir()}/dump.sql'''
-    ])
+    cmd = " ".join(
+        [
+            "pg_dump -Fp",
+            f'''-d "{db_url(USER, "src")}"''',
+            f"""-f {docker_dir()}/dump.sql""",
+        ]
+    )
     if where_am_i() != "docker":
         cmd = f"docker exec -w {docker_dir()} pgai {cmd}"
     subprocess.run(cmd, check=True, shell=True, env=os.environ, cwd=str(host_dir()))
@@ -68,39 +77,45 @@ def restore_db() -> None:
     with psycopg.connect(db_url(user=USER, dbname="dst")) as con:
         with con.cursor() as cur:
             cur.execute("create extension ai cascade")
-    cmd = " ".join([
-        "psql",
-        f'''-d "{db_url(USER, "dst")}"''',
-        "-v VERBOSITY=verbose",
-        f"-f {docker_dir()}/dump.sql",
-    ])
+    cmd = " ".join(
+        [
+            "psql",
+            f'''-d "{db_url(USER, "dst")}"''',
+            "-v VERBOSITY=verbose",
+            f"-f {docker_dir()}/dump.sql",
+        ]
+    )
     if where_am_i() != "docker":
         cmd = f"docker exec -w {docker_dir()} pgai {cmd}"
     subprocess.run(cmd, check=True, shell=True, env=os.environ, cwd=str(host_dir()))
 
 
 def snapshot_db(dbname: str) -> None:
     host_dir().joinpath(f"{dbname}.snapshot").unlink(missing_ok=True)
-    cmd = " ".join([
-        "psql",
-        f'''-d "{db_url("postgres", dbname)}"''',
-        "-v ON_ERROR_STOP=1",
-        "-X",
-        f"-o {docker_dir()}/{dbname}.snapshot",
-        f"-f {docker_dir()}/snapshot.sql",
-    ])
+    cmd = " ".join(
+        [
+            "psql",
+            f'''-d "{db_url("postgres", dbname)}"''',
+            "-v ON_ERROR_STOP=1",
+            "-X",
+            f"-o {docker_dir()}/{dbname}.snapshot",
+            f"-f {docker_dir()}/snapshot.sql",
+        ]
+    )
     if where_am_i() != "docker":
         cmd = f"docker exec -w {docker_dir()} pgai {cmd}"
     subprocess.run(cmd, check=True, shell=True, env=os.environ, cwd=str(host_dir()))
 
 
 def init_src() -> None:
-    cmd = " ".join([
-        "psql",
-        f'''-d "{db_url(USER, "src")}"''',
-        "-v ON_ERROR_STOP=1",
-        f"-f {docker_dir()}/init.sql",
-    ])
+    cmd = " ".join(
+        [
+            "psql",
+            f'''-d "{db_url(USER, "src")}"''',
+            "-v ON_ERROR_STOP=1",
+            f"-f {docker_dir()}/init.sql",
+        ]
+    )
     if where_am_i() != "docker":
         cmd = f"docker exec -w {docker_dir()} pgai {cmd}"
     subprocess.run(cmd, check=True, shell=True, env=os.environ, cwd=str(host_dir()))
@@ -112,12 +127,14 @@ def read_file(filename: str) -> str:
 
 
 def after_dst() -> None:
-    cmd = " ".join([
-        "psql",
-        f'''-d "{db_url(USER, "dst")}"''',
-        "-v ON_ERROR_STOP=1",
-        f"-f {docker_dir()}/after.sql",
-    ])
+    cmd = " ".join(
+        [
+            "psql",
+            f'''-d "{db_url(USER, "dst")}"''',
+            "-v ON_ERROR_STOP=1",
+            f"-f {docker_dir()}/after.sql",
+        ]
+    )
     if where_am_i() != "docker":
         cmd = f"docker exec -w {docker_dir()} pgai {cmd}"
     subprocess.run(cmd, check=True, shell=True, env=os.environ, cwd=str(host_dir()))

projects/extension/tests/privileges/test_privileges.py

@@ -35,13 +35,15 @@ def read_file(filename: str) -> str:
 
 
 def psql_file(user, dbname, file: str) -> None:
-    cmd = " ".join([
-        "psql",
-        f'''-d "{db_url(user, dbname)}"''',
-        "-v ON_ERROR_STOP=1",
-        "-X",
-        f"-f {docker_dir()}/{file}",
-    ])
+    cmd = " ".join(
+        [
+            "psql",
+            f'''-d "{db_url(user, dbname)}"''',
+            "-v ON_ERROR_STOP=1",
+            "-X",
+            f"-f {docker_dir()}/{file}",
+        ]
+    )
     if where_am_i() != "docker":
         cmd = f"docker exec -w {docker_dir()} pgai {cmd}"
     subprocess.run(cmd, check=True, shell=True, env=os.environ, cwd=str(host_dir()))
@@ -82,4 +84,3 @@ def test_function_privileges():
 
 def test_jill_privileges():
     psql_file("jill", "privs", "jill.sql")
-

projects/extension/tests/vectorizer/server.py

@@ -13,8 +13,8 @@
 
 
 class ChunkingCharacterTextSplitter(BaseModel):
-    implementation: Literal['character_text_splitter']
-    config_type: Literal['chunking']
+    implementation: Literal["character_text_splitter"]
+    config_type: Literal["chunking"]
     chunk_column: str
     chunk_size: int
     chunk_overlap: int
@@ -23,8 +23,8 @@ class ChunkingCharacterTextSplitter(BaseModel):
 
 
 class ChunkingRecursiveCharacterTextSplitter(BaseModel):
-    implementation: Literal['recursive_character_text_splitter']
-    config_type: Literal['chunking']
+    implementation: Literal["recursive_character_text_splitter"]
+    config_type: Literal["chunking"]
     chunk_column: str
     chunk_size: int
     chunk_overlap: int
@@ -33,30 +33,30 @@ class ChunkingRecursiveCharacterTextSplitter(BaseModel):
 
 
 class EmbeddingOpenAI(BaseModel):
-    implementation: Literal['openai']
-    config_type: Literal['embedding']
+    implementation: Literal["openai"]
+    config_type: Literal["embedding"]
     model: str
     dimensions: int
     user: str | None = None
-    api_key_name: str | None = 'OPENAI_API_KEY'
+    api_key_name: str | None = "OPENAI_API_KEY"
 
 
 class FormattingPythonTemplate(BaseModel):
-    implementation: Literal['python_template']
-    config_type: Literal['formatting']
+    implementation: Literal["python_template"]
+    config_type: Literal["formatting"]
     template: str
 
 
 class IndexingNone(BaseModel):
-    implementation: Literal['none']
-    config_type: Literal['indexing']
+    implementation: Literal["none"]
+    config_type: Literal["indexing"]
 
 
 class IndexingDiskANN(BaseModel):
-    implementation: Literal['diskann']
-    config_type: Literal['indexing']
+    implementation: Literal["diskann"]
+    config_type: Literal["indexing"]
     min_rows: int
-    storage_layout: Literal['memory_optimized'] | Literal['plain'] | None = None
+    storage_layout: Literal["memory_optimized"] | Literal["plain"] | None = None
     num_neighbors: int | None = None
     search_list_size: int | None = None
     max_alpha: float | None = None
@@ -65,54 +65,67 @@ class IndexingDiskANN(BaseModel):
 
 
 class IndexingHNSW(BaseModel):
-    implementation: Literal['hnsw']
-    config_type: Literal['indexing']
+    implementation: Literal["hnsw"]
+    config_type: Literal["indexing"]
     min_rows: int
-    opclass: Literal['vector_ip_ops'] | Literal['vector_cosine_ops'] | Literal['vector_l1_ops'] | None = None
+    opclass: (
+        Literal["vector_ip_ops"]
+        | Literal["vector_cosine_ops"]
+        | Literal["vector_l1_ops"]
+        | None
+    ) = None
     m: int | None = None
     ef_construction: int | None = None
 
 
 class SchedulingNone(BaseModel):
-    implementation: Literal['none']
-    config_type: Literal['scheduling']
+    implementation: Literal["none"]
+    config_type: Literal["scheduling"]
 
 
 class SchedulingPgCron(BaseModel):
-    implementation: Literal['pg_cron']
-    config_type: Literal['scheduling']
+    implementation: Literal["pg_cron"]
+    config_type: Literal["scheduling"]
     schedule: str
 
 
 class SchedulingTimescaledb(BaseModel):
-    implementation: Literal['timescaledb']
-    config_type: Literal['scheduling']
+    implementation: Literal["timescaledb"]
+    config_type: Literal["scheduling"]
     schedule_interval: str | None = None
     initial_start: datetime | None = None
     fixed_schedule: bool | None = None
     timezone: str | None = None
 
 
 class ProcessingNone(BaseModel):
-    implementation: Literal['none']
-    config_type: Literal['processing']
+    implementation: Literal["none"]
+    config_type: Literal["processing"]
 
 
 class ProcessingCloudFunctions(BaseModel):
-    implementation: Literal['cloud_functions']
-    config_type: Literal['processing']
+    implementation: Literal["cloud_functions"]
+    config_type: Literal["processing"]
     batch_size: int | None = None
     concurrency: int | None = None
 
 
 class Config(BaseModel):
     version: str
-    indexing: Union[IndexingNone, IndexingDiskANN, IndexingHNSW] = Field(..., discriminator='implementation')
+    indexing: Union[IndexingNone, IndexingDiskANN, IndexingHNSW] = Field(
+        ..., discriminator="implementation"
+    )
     formatting: FormattingPythonTemplate
     embedding: EmbeddingOpenAI
-    scheduling: Union[SchedulingNone, SchedulingPgCron, SchedulingTimescaledb] = Field(..., discriminator='implementation')
-    chunking: Union[ChunkingCharacterTextSplitter, ChunkingRecursiveCharacterTextSplitter] = Field(..., discriminator='implementation')
-    processing: Union[ProcessingNone, ProcessingCloudFunctions] = Field(..., discriminator='implementation')
+    scheduling: Union[SchedulingNone, SchedulingPgCron, SchedulingTimescaledb] = Field(
+        ..., discriminator="implementation"
+    )
+    chunking: Union[
+        ChunkingCharacterTextSplitter, ChunkingRecursiveCharacterTextSplitter
+    ] = Field(..., discriminator="implementation")
+    processing: Union[ProcessingNone, ProcessingCloudFunctions] = Field(
+        ..., discriminator="implementation"
+    )
 
 
 class PrimaryKeyColumn(BaseModel):
@@ -162,4 +175,3 @@ async def execute_vectorizer(vectorizer: Vectorizer):
     deleted = vectorize(vectorizer)
     print(f"queue emptied: {deleted} rows deleted")
     return {"id": vectorizer.id}
-