Split extension into an extension and a library package #580

cevian · 2025-03-21T21:12:13Z

Purpose

This PR separates the codebase into two distinct components:

The Extension: Focused specifically on model-calling functionality
The pgai Python Library: Houses database functionality installed via Python scripts, including
vectorizer and text2sql functionality

The key advantage: Portability. The pgai Python library can now be used on any PostgreSQL
installation, including cloud platforms that don't support the pgai extensions.

Architecture

The pgai library installs database objects (tables, schemas, functions) through the pgai install
function using scripts compiled from the projects/pgai/db directory. The migration management follows
the extension's proven approach:

Idempotent SQL: Safe to run on each installation, primarily containing function definitions
Incremental SQL: Ordered migrations that evolve database objects sequentially, tracked via the
ai.migration_app table

The approach of using the ai.migration_app table allows us to maintains a single upgrade script compatible with any previous version, eliminating the need for version-specific upgrade paths.

Migration Strategy

To transition from the previous monolithic extension:

Code Relocation: Vectorizer SQL code (both idempotent and incremental) moved from extension to
pgai/db. Since the vectorizer code was mostly in their own files, these files moved mostly unchanged.
Divest Process: Created a script that cleanly dissociates vectorizer-related objects from the
extension while preserving their functionality
Migration Tracking: Recorded previously extension-applied vectorizer-related incremental files in
ai.migration_app

After upgrading to the next extension version, vectorizer components remain intact but are no longer managed by the extension. The extension does not upgrade the vectorizer components before releasing them, since it now lack the required code. Instead, the pgai library can then independently upgrade these components regardless of which extension version previously created them.

projects/extension/build.py

projects/extension/sql/idempotent/012-vectorizer-int.sql

projects/extension/Dockerfile

projects/extension/sql/head.sql

projects/extension/tests/dump_restore/after.sql

projects/pgai/db/justfile

jgpruitt · 2025-03-24T16:20:01Z

projects/pgai/db/tests/vectorizer/test_vectorizer.py

@@ -177,6 +184,10 @@ def psql_cmd(cmd: str) -> str:


 def test_vectorizer_timescaledb():
+    with psycopg.connect(db_url("test")) as con:
+        with con.cursor() as cur:
+            cur.execute("create extension ai cascade")


i don't understand why we're installing the extension?

because the thing that sends the ping to the lamda is still in the extension so on our cloud there is a dependency

I added logic to the install() function to set up the ai extension if executor_url is set

jgpruitt · 2025-03-24T16:24:29Z

projects/pgai/pgai/_install/install.py

Shouldn't we take an advisory lock before attempting install/upgrade? bail if we don't get it.

jgpruitt · 2025-03-24T16:28:58Z

projects/pgai/db/sql/head.sql

+);
+
+insert into ai.app_version ("name", version)
+values ('ai', '__version__') on conflict ("name") do update set version = excluded.version;


Should we check the version? If the version we are trying to install/upgrade is already present, abort? If you're running N apps in a load balancer, we probably want the first to win and the rest to skip.

projects/pgai/db/sql/head.sql

JamesGuthrie · 2025-03-25T07:36:38Z

projects/pgai/.python-version

Why was this removed?

I think this was committed by mistake in the first place? Seems inappropriate on a library that can run on many python versions

JamesGuthrie · 2025-03-25T07:41:16Z

projects/pgai/db/README.md

+-> leaving ai.execute_vectorizer in extension
+->  ai._vectorizer_job which calls ai.execute_vectorizer is moved to dbapp
+
+-> creating the index still needs to be handled (probably through the worker)


This doesn't work for lambda-based workers.

For lamba stuff it's done in the timescaledb job. That actually already works

projects/extension/sql/incremental/020-divest.sql

projects/pgai/db/sql/flag.sql

projects/pgai/db/sql/head.sql

projects/pgai/pgai/_install/install.py

projects/pgai/pgai/cli.py

projects/pgai/pgai/vectorizer/features/features.py

projects/pgai/tests/vectorizer/cli/conftest.py

projects/extension/tests/privileges/test_privileges.py

projects/extension/tests/upgrade/test_upgrade.py

projects/pgai/db/sql/incremental/frozen.txt

projects/extension/tests/upgrade/test_upgrade.py

jgpruitt · 2025-03-31T18:46:44Z

projects/pgai/db/build.py

+def error_if_pre_release() -> None:
+    # Note: released versions always have the output sql file
+    # commited into the repository.
+    output_file = output_sql_file()
+    command = (
+        "just ext build-install"
+        if "ROOT_JUSTFILE" in os.environ
+        else "just build-install"
+        if "PROJECT_JUSTFILE" in os.environ
+        else "python3 build.py build-install"
+    )
+    if not Path(output_file).exists():
+        print(
+            textwrap.dedent(f"""
+                WARNING: You're trying to install a pre-release version of pgai.
+                This is not supported, and there is no upgrade path.
+
+                Instead, install an official release from https://github.com/timescale/pgai/releases.
+
+                If you are certain that you want to install a pre-release version, run:
+                    `{command}`
+            """)
+        )
+        exit(1)


I don't think this is quite right. Also, I'm not sure it's necessary. People won't be using the just commands to install the SQL anymore. Maybe this check needs to go into the Python library?

yeah this was dead code. I removed it and added a check in the python code.

projects/pgai/pgai/_install/install.py

projects/pgai/db/tests/vectorizer/test_vectorizer.py

projects/pgai/pgai/cli.py

jgpruitt · 2025-03-31T19:34:01Z

projects/pgai/pgai/cli.py

+    default="postgres://postgres@localhost:5432/postgres",
+    show_default=True,
+)
+def install(db_url: str) -> None:


Is upgrading the same as installing over an existing installation?

changes made

JamesGuthrie

A few comments, but approving to unblock.

JamesGuthrie · 2025-04-01T09:12:53Z

projects/extension/sql/incremental/020-divest.sql

+        raise notice '%', _sql;
+        execute _sql;
+
+        if _rec.relname != 'vectorizer_id_seq' THEN


Why do we not change the ownership of vectorizer_id_seq?

added a code comment. The reason is it's linked to the vectorizer table so you change the owner of the table instead.

JamesGuthrie · 2025-04-01T09:38:15Z

.github/workflows/ci.yml

These changes aren't tested by CI because we use pull_request_target to build it. I attempted to run just pgai db ci locally, and it didn't pass, so I don't have confidence that these tests will succeed when this merges.

I (accidentally) ran them by creating a PR to this branch and they are red: https://github.com/timescale/pgai/actions/runs/14193958551/job/39764748178?pr=598

run manually through workflow_dispatch. everything passes now https://github.com/timescale/pgai/actions/runs/14205814536

JamesGuthrie · 2025-04-01T09:54:17Z

projects/pgai/db/build.py

+            [
+                "uv run --no-project pgspot --ignore-lang=plpython3u",
+                '--proc-without-search-path "ai._vectorizer_job(job_id integer,config pg_catalog.jsonb)"',
+                f"{output_sql_file()}",


pushed a commit with this fix.

JamesGuthrie · 2025-04-01T10:00:53Z

projects/extension/Dockerfile

+USER root
+
+RUN pip install --break-system-packages uv==0.6.3
+RUN mkdir /py/ && uv venv --directory /py/


This is missing correctly setting up dependencies in the container.

Suggested change

RUN mkdir /py/ && uv venv --directory /py/

COPY pyproject.toml uv.lock /py/

RUN uv sync --directory /py --no-install-project --only-dev --frozen

yeah I do that in the just command docker-sync now. Let's review after merge.

JamesGuthrie · 2025-04-01T10:01:45Z

projects/pgai/db/README.md

+just docker-shell
+
+//inside shell 1
+uv pip install fastapi[standard]


With the suggested change to the dockerfile, this becomes redundant.

Suggested change

uv pip install fastapi[standard]

yeah these were just outdated docs docker-sync takes care of this

JamesGuthrie · 2025-04-01T10:03:12Z

projects/pgai/db/README.md

+uv sync --active
+cd ../extension && just build && just install && cd ../pgai/db


We should eliminate these extra manual steps required to run the tests.

yeah outdate docs this is already not necessary. fixing doc

JamesGuthrie · 2025-04-01T10:32:17Z

projects/pgai/db/sql/head.sql

+, installed_at timestamptz not null default pg_catalog.clock_timestamp()
+);
+
+--check if the app has already been installed, error if so


Why error? Because it's the only way to stop with the rest of the installation, or because we actually see this as being erroneous?

yes because its the only way to stop the install

projects/extension/tests/upgrade/test_upgrade.py

projects/pgai/db/build.py

cevian temporarily deployed to internal-contributors March 21, 2025 21:12 — with GitHub Actions Inactive

cevian temporarily deployed to internal-contributors March 23, 2025 18:40 — with GitHub Actions Inactive

jgpruitt requested changes Mar 24, 2025

View reviewed changes

JamesGuthrie reviewed Mar 25, 2025

View reviewed changes

cevian force-pushed the mat/repackage branch from c7d5163 to 27503d6 Compare March 28, 2025 16:55

cevian temporarily deployed to internal-contributors March 28, 2025 16:56 — with GitHub Actions Inactive

cevian marked this pull request as ready for review March 28, 2025 16:57

cevian requested a review from a team as a code owner March 28, 2025 16:57

feat: new packaging - move files into new project

90ba526

cevian force-pushed the mat/repackage branch from 27503d6 to d47c143 Compare March 28, 2025 17:20

cevian temporarily deployed to internal-contributors March 28, 2025 17:21 — with GitHub Actions Inactive

cevian temporarily deployed to internal-contributors March 28, 2025 21:53 — with GitHub Actions Inactive

cevian force-pushed the mat/repackage branch from 1cb1976 to e84f500 Compare March 28, 2025 21:59

cevian temporarily deployed to internal-contributors March 28, 2025 21:59 — with GitHub Actions Inactive

cevian force-pushed the mat/repackage branch from e84f500 to 82b3584 Compare March 28, 2025 22:00

cevian temporarily deployed to internal-contributors March 28, 2025 22:01 — with GitHub Actions Inactive

cevian force-pushed the mat/repackage branch from 82b3584 to 1f24ad9 Compare March 30, 2025 17:52

cevian temporarily deployed to internal-contributors March 30, 2025 17:52 — with GitHub Actions Inactive

cevian temporarily deployed to internal-contributors March 30, 2025 18:10 — with GitHub Actions Inactive

cevian requested review from jgpruitt and JamesGuthrie March 30, 2025 18:12

cevian temporarily deployed to internal-contributors March 30, 2025 18:25 — with GitHub Actions Inactive

cevian temporarily deployed to internal-contributors March 30, 2025 18:26 — with GitHub Actions Inactive

cevian temporarily deployed to internal-contributors March 30, 2025 18:36 — with GitHub Actions Inactive

cevian temporarily deployed to internal-contributors March 30, 2025 18:37 — with GitHub Actions Inactive

cevian force-pushed the mat/repackage branch from d7ffbac to 4777b76 Compare March 30, 2025 18:55

cevian temporarily deployed to internal-contributors March 30, 2025 18:55 — with GitHub Actions Inactive

cevian temporarily deployed to internal-contributors March 30, 2025 18:57 — with GitHub Actions Inactive

cevian temporarily deployed to internal-contributors March 30, 2025 19:30 — with GitHub Actions Inactive

fix: addressing first round of comments

36e7344

fix: change _app -> pgai_lib

a29a78c

cevian force-pushed the mat/repackage branch from 5cbd04f to a29a78c Compare March 31, 2025 01:03

cevian temporarily deployed to internal-contributors March 31, 2025 01:03 — with GitHub Actions Inactive

jgpruitt requested changes Mar 31, 2025

View reviewed changes

fix: john's comments

9913757

cevian temporarily deployed to internal-contributors April 1, 2025 01:20 — with GitHub Actions Inactive

fix: s/output_sql_file/lib_sql_file/

6e924be

JamesGuthrie temporarily deployed to internal-contributors April 1, 2025 09:54 — with GitHub Actions Inactive

JamesGuthrie approved these changes Apr 1, 2025

View reviewed changes

cevian added 2 commits April 1, 2025 13:33

fix: db tests

c39928a

fix: move privilege tests over

e9e9f87

cevian temporarily deployed to internal-contributors April 1, 2025 19:04 — with GitHub Actions Inactive

fix: rename if_not_exist->strict, update logic

a954a47

cevian temporarily deployed to internal-contributors April 1, 2025 19:37 — with GitHub Actions Inactive

cevian added 2 commits April 1, 2025 16:11

fix: james comments

4d342d9

fix: format and lint

40af350

cevian temporarily deployed to internal-contributors April 1, 2025 20:13 — with GitHub Actions Inactive

cevian temporarily deployed to internal-contributors April 1, 2025 20:14 — with GitHub Actions Inactive

fix: tests that require ai extension

c0fff7f

cevian temporarily deployed to internal-contributors April 1, 2025 20:53 — with GitHub Actions Inactive

cevian requested a review from jgpruitt April 1, 2025 20:54

test: make pgai install work from host or docker

33b91fb

jgpruitt temporarily deployed to internal-contributors April 1, 2025 22:15 — with GitHub Actions Inactive

jgpruitt approved these changes Apr 2, 2025

View reviewed changes

projects/extension/tests/upgrade/test_upgrade.py Outdated Show resolved Hide resolved

projects/pgai/db/build.py Show resolved Hide resolved

fix: small fix

2d57ed6

cevian temporarily deployed to internal-contributors April 2, 2025 15:29 — with GitHub Actions Inactive

cevian merged commit 3fe83c6 into main Apr 2, 2025
13 checks passed

cevian deleted the mat/repackage branch April 2, 2025 16:04

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Split extension into an extension and a library package #580

Split extension into an extension and a library package #580

cevian commented Mar 21, 2025 •

edited

Loading

jgpruitt Mar 24, 2025

cevian Mar 26, 2025

cevian Mar 26, 2025

jgpruitt Mar 24, 2025

jgpruitt Mar 24, 2025

JamesGuthrie Mar 25, 2025

cevian Mar 25, 2025

JamesGuthrie Mar 25, 2025

cevian Mar 25, 2025

jgpruitt Mar 31, 2025

cevian Apr 1, 2025

jgpruitt Mar 31, 2025

cevian Apr 1, 2025

JamesGuthrie left a comment

JamesGuthrie Apr 1, 2025

cevian Apr 1, 2025

JamesGuthrie Apr 1, 2025

Askir Apr 1, 2025 •

edited

Loading

cevian Apr 1, 2025

JamesGuthrie Apr 1, 2025

JamesGuthrie Apr 1, 2025

cevian Apr 1, 2025

JamesGuthrie Apr 1, 2025

cevian Apr 1, 2025

JamesGuthrie Apr 1, 2025

cevian Apr 1, 2025

JamesGuthrie Apr 1, 2025

cevian Apr 1, 2025

	RUN mkdir /py/ && uv venv --directory /py/
	COPY pyproject.toml uv.lock /py/
	RUN uv sync --directory /py --no-install-project --only-dev --frozen

		uv sync --active
		cd ../extension && just build && just install && cd ../pgai/db

Split extension into an extension and a library package #580

Split extension into an extension and a library package #580

Conversation

cevian commented Mar 21, 2025 • edited Loading

Purpose

Architecture

Migration Strategy

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JamesGuthrie left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Askir Apr 1, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cevian commented Mar 21, 2025 •

edited

Loading

Askir Apr 1, 2025 •

edited

Loading