From 564474f51ac5b756c1e87a490fa20873f2895372 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Fri, 30 Oct 2026 07:35:35 -0700 Subject: [PATCH] feat: use SQL for remote function lifecycle Remote function registration, exact-version lookup, listing, and drop now execute SQL through the Flight client. Registration waits for the durable image build and catalog publication before yielding its typed `FunctionVersion`. SQL also covers source-inclusive description and pinned metadata updates. Registration and drop handles use `job_type="query"`; the single `job.id` is the SQL query UUID and can be inspected with `describe_query()` on the submitting connection. Waiting drains the SQL result once and caches completion for repeated waits. The drop boolean reports catalog removal; physical cleanup continues separately. Query handles have no registry history and cannot be reopened with `open_job()`. --- docs/web-tests/py/test_connection.py | 18 + docs/web-tests/py/test_function.py | 20 +- docs/web-tests/ts/connection.test.ts | 17 +- docs/web/basics/jobs.mdx | 6 + docs/web/management/function.mdx | 13 +- docs/web/snippets/connection.mdx | 8 +- docs/web/snippets/function.mdx | 2 +- docs/web/sql/ddl/function.mdx | 53 +- python/python/lancedb/db.py | 33 +- .../tests/test_first_class_function_slice2.py | 310 +--------- python/python/tests/test_listing.py | 2 +- rust/lancedb/src/connection.rs | 16 +- rust/lancedb/src/database.rs | 7 +- rust/lancedb/src/remote/db.rs | 546 ++---------------- rust/lancedb/src/remote/sql.rs | 2 + rust/lancedb/src/remote/sql/function.rs | 192 ++++++ rust/lancedb/src/remote/sql/objects.rs | 72 +++ rust/lancedb/src/remote/sql_test.rs | 164 ++++++ 18 files changed, 607 insertions(+), 874 deletions(-) create mode 100644 rust/lancedb/src/remote/sql/function.rs create mode 100644 rust/lancedb/src/remote/sql/objects.rs diff --git a/docs/web-tests/py/test_connection.py b/docs/web-tests/py/test_connection.py index f5c1f2b22..b7a6256ca 100644 --- a/docs/web-tests/py/test_connection.py +++ b/docs/web-tests/py/test_connection.py @@ -41,6 +41,7 @@ def connect_enterprise_quickstart_config(): api_key=os.environ["LANCEDB_API_KEY"], region=os.environ.get("LANCEDB_REGION", "us-east-1"), host_override=os.environ["LANCEDB_HOST_OVERRIDE"], # https://your-endpoint + sql_host_override=os.environ["LANCEDB_SQL_HOST_OVERRIDE"], ) # --8<-- [end:connect_enterprise_quickstart] return db @@ -60,6 +61,9 @@ def test_connect_enterprise_quickstart(monkeypatch): monkeypatch.setenv("LANCEDB_API_KEY", "your-api-key") monkeypatch.delenv("LANCEDB_REGION", raising=False) monkeypatch.setenv("LANCEDB_HOST_OVERRIDE", "https://your-endpoint") + monkeypatch.setenv( + "LANCEDB_SQL_HOST_OVERRIDE", "grpc+tls://your-sql-endpoint:10026" + ) assert connect_enterprise_quickstart_config() is not None assert captured == { @@ -67,6 +71,7 @@ def test_connect_enterprise_quickstart(monkeypatch): "api_key": "your-api-key", "region": "us-east-1", "host_override": "https://your-endpoint", + "sql_host_override": "grpc+tls://your-sql-endpoint:10026", } @@ -79,6 +84,16 @@ def test_connect_enterprise_quickstart_requires_an_endpoint(monkeypatch): connect_enterprise_quickstart_config() +def test_connect_enterprise_quickstart_requires_a_sql_endpoint(monkeypatch): + monkeypatch.setenv("LANCEDB_URI", "db://your-database") + monkeypatch.setenv("LANCEDB_API_KEY", "your-api-key") + monkeypatch.setenv("LANCEDB_HOST_OVERRIDE", "https://your-endpoint") + monkeypatch.delenv("LANCEDB_SQL_HOST_OVERRIDE", raising=False) + + with pytest.raises(KeyError, match="LANCEDB_SQL_HOST_OVERRIDE"): + connect_enterprise_quickstart_config() + + def connect_object_storage_config(): # --8<-- [start:connect_object_storage] import lancedb @@ -169,6 +184,7 @@ def test_namespace_admin_ops(tmp_path, monkeypatch, capsys): ) assert db.list_namespaces().namespaces == [] + async def connect_object_storage_config_async(): # --8<-- [start:connect_object_storage_async] import lancedb @@ -199,6 +215,7 @@ def test_connect_enterprise_placeholders(monkeypatch): api_key="your-api-key", region="us-east-1", host_override="https://your-endpoint", + sql_host_override="grpc+tls://your-sql-endpoint:10026", ) # --8<-- [end:connect_enterprise_placeholders] assert db is not None @@ -207,4 +224,5 @@ def test_connect_enterprise_placeholders(monkeypatch): "api_key": "your-api-key", "region": "us-east-1", "host_override": "https://your-endpoint", + "sql_host_override": "grpc+tls://your-sql-endpoint:10026", } diff --git a/docs/web-tests/py/test_function.py b/docs/web-tests/py/test_function.py index a89424316..297ac718b 100644 --- a/docs/web-tests/py/test_function.py +++ b/docs/web-tests/py/test_function.py @@ -6,8 +6,9 @@ Registering a Function and filling its column need a LanceDB Enterprise deployment, so `test_function_enterprise` runs only when LANCEDB_DOCS_TEST_DEPLOYMENT=1 asks for one, with the connection variables the -Enterprise examples already use: LANCEDB_URI, LANCEDB_API_KEY and -LANCEDB_HOST_OVERRIDE, plus LANCEDB_REGION if the deployment needs one. The +Enterprise examples already use: LANCEDB_URI, LANCEDB_API_KEY, +LANCEDB_HOST_OVERRIDE and LANCEDB_SQL_HOST_OVERRIDE, plus LANCEDB_REGION if the +deployment needs one. The walkthrough calls its table `function_demo` and its Function `double`; the test creates both under names unique to its run, never replaces an existing object, and drops only what it created. @@ -91,7 +92,12 @@ def registered_version(definition) -> FunctionVersion: return FunctionVersion.from_json(json.dumps(reply)) -DEPLOYMENT = ("LANCEDB_URI", "LANCEDB_HOST_OVERRIDE", "LANCEDB_API_KEY") +DEPLOYMENT = ( + "LANCEDB_URI", + "LANCEDB_HOST_OVERRIDE", + "LANCEDB_SQL_HOST_OVERRIDE", + "LANCEDB_API_KEY", +) # The connection settings alone never send this test to a deployment. OPT_IN = "LANCEDB_DOCS_TEST_DEPLOYMENT" @@ -181,7 +187,10 @@ def test_function_version_binds_a_column(tmp_path): # The binding is accepted by add_columns, which a local table refuses: # a Function column needs LanceDB Cloud or Enterprise. table = lancedb.connect(tmp_path).create_table("t", [{"value": 1.0}]) - with pytest.raises(NotImplementedError, match="Cloud and Enterprise"): + with pytest.raises( + NotImplementedError, + match="Function columns are supported only in LanceDB Distributed", + ): table.add_columns({"doubled": application}) @@ -223,6 +232,7 @@ def test_function_enterprise(monkeypatch): api_key=os.environ["LANCEDB_API_KEY"], region=os.environ.get("LANCEDB_REGION", "us-east-1"), host_override=os.environ["LANCEDB_HOST_OVERRIDE"], + sql_host_override=os.environ["LANCEDB_SQL_HOST_OVERRIDE"], ) # --8<-- [end:function_connect] @@ -298,6 +308,7 @@ class ReachedConnection(Exception): SETTINGS = { "LANCEDB_URI": "db://docs-test", "LANCEDB_HOST_OVERRIDE": "https://deployment.invalid", + "LANCEDB_SQL_HOST_OVERRIDE": "grpc+tls://sql.deployment.invalid:10026", "LANCEDB_API_KEY": "docs-test-key", } @@ -375,5 +386,6 @@ def test_function_enterprise_connects_when_asked(monkeypatch, connect_attempts, "api_key": "docs-test-key", "region": region or "us-east-1", "host_override": "https://deployment.invalid", + "sql_host_override": "grpc+tls://sql.deployment.invalid:10026", } ] diff --git a/docs/web-tests/ts/connection.test.ts b/docs/web-tests/ts/connection.test.ts index 651be7d36..8cab34c26 100644 --- a/docs/web-tests/ts/connection.test.ts +++ b/docs/web-tests/ts/connection.test.ts @@ -93,6 +93,12 @@ async function connectObjectStorageExample() { async function namespaceTableOpsExample() { // --8<-- [start:namespace_table_ops] + async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; + } + const db = await lancedb.connectNamespace("dir", { root: "./local_lancedb" }); // Create namespace tree: prod/search and prod/recommendations @@ -124,6 +130,12 @@ async function namespaceTableOpsExample() { async function namespaceAdminOpsExample() { // --8<-- [start:namespace_admin_ops] + async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; + } + const db = await lancedb.connectNamespace("dir", { root: "./local_lancedb" }); const namespace = ["prod", "search"]; @@ -131,7 +143,10 @@ async function namespaceAdminOpsExample() { await db.createNamespace(["prod", "search"]); const childNamespaces = await collectNames(db.listNamespaces(["prod"])); - console.log(`Child namespaces under ${JSON.stringify(namespace)}:`, childNamespaces); + console.log( + `Child namespaces under ${JSON.stringify(namespace)}:`, + childNamespaces, + ); // Child namespaces under ["prod","search"]: [ 'search' ] const metadata = await db.describeNamespace(["prod", "search"]); diff --git a/docs/web/basics/jobs.mdx b/docs/web/basics/jobs.mdx index a15f10fee..9f4440de0 100644 --- a/docs/web/basics/jobs.mdx +++ b/docs/web/basics/jobs.mdx @@ -119,6 +119,12 @@ and `cancel()`. Synchronous creation, refresh, and drop wait on the same SQL execution before returning. A completed DROP means the catalog binding has been removed; physical cleanup continues in the background. +Remote function registration also returns a SQL query job. The SQL operation +submits the durable image build and waits for catalog publication before returning +the immutable FunctionVersion. `drop_function_async()` returns the catalog-removal +boolean together with a query job; waiting on that handle confirms removal while +physical cleanup continues separately. + The existing `job.job_type` identifies query-backed handles as `"query"` from submission, and `job.id` contains the connection-scoped SQL query UUID string. Use `describe_query(job.id)` on the submitting connection to inspect the query. diff --git a/docs/web/management/function.mdx b/docs/web/management/function.mdx index 415de7613..b7fbb504b 100644 --- a/docs/web/management/function.mdx +++ b/docs/web/management/function.mdx @@ -91,7 +91,9 @@ deployment's support before using another computed column as an input. ## Register it {#register} -Registration returns a job; the immutable version exists once the job succeeds. +Registration returns a query-backed job; the immutable version exists once the +SQL command finishes and the job succeeds. Its `job_type` is `query`, and its +single `id` is the SQL query UUID. Repeated waits reuse the completed result. {FunctionRegister} @@ -101,8 +103,9 @@ Registration returns a job; the immutable version exists once the job succeeds. with `db.get_function(name, version=...)`. -**Enterprise job type:** `create_function`. Registration returns a job because -the function is baked into an image before the version becomes usable. +The SQL command submits an internal `create_function` image-build job and +waits for publication. The SDK returns the SQL query handle. Query handles +have no registry history and cannot be reopened with `open_job()`. ## Try it in SQL {#try-sql} @@ -565,7 +568,9 @@ itself is: as source, snapshotted at registration. it is bound to; `version` must be a version that exists. It returns `False` if the name was already gone. Deleting the object's content can finish after the call returns, as a [background job](/maintenance/background-jobs). -`drop_function_async` returns that job with the result, so you can wait for it. +`drop_function_async` returns `(dropped, job)`, where `job` is a query-backed +handle. Waiting confirms catalog +removal; physical cleanup can continue afterward. A deployment refuses the drop unless it also allows changes to a function's lifecycle, which is a separate setting, off by default even where Functions are diff --git a/docs/web/snippets/connection.mdx b/docs/web/snippets/connection.mdx index c145c1fb6..fd8164ecf 100644 --- a/docs/web/snippets/connection.mdx +++ b/docs/web/snippets/connection.mdx @@ -4,9 +4,9 @@ export const PyConnect = "import lancedb\n\nuri = \"ex_lancedb\"\ndb = lancedb.c export const PyConnectAsync = "import lancedb\n\nuri = \"ex_lancedb\"\nasync_db = await lancedb.connect_async(uri)\n"; -export const PyConnectEnterprisePlaceholders = "import lancedb\n\ndb = lancedb.connect(\n uri=\"db://your-database\",\n api_key=\"your-api-key\",\n region=\"us-east-1\",\n host_override=\"https://your-endpoint\",\n)\n"; +export const PyConnectEnterprisePlaceholders = "import lancedb\n\ndb = lancedb.connect(\n uri=\"db://your-database\",\n api_key=\"your-api-key\",\n region=\"us-east-1\",\n host_override=\"https://your-endpoint\",\n sql_host_override=\"grpc+tls://your-sql-endpoint:10026\",\n)\n"; -export const PyConnectEnterpriseQuickstart = "import os\n\nimport lancedb\n\ndb = lancedb.connect(\n uri=os.environ[\"LANCEDB_URI\"], # db://your-database\n api_key=os.environ[\"LANCEDB_API_KEY\"],\n region=os.environ.get(\"LANCEDB_REGION\", \"us-east-1\"),\n host_override=os.environ[\"LANCEDB_HOST_OVERRIDE\"], # https://your-endpoint\n)\n"; +export const PyConnectEnterpriseQuickstart = "import os\n\nimport lancedb\n\ndb = lancedb.connect(\n uri=os.environ[\"LANCEDB_URI\"], # db://your-database\n api_key=os.environ[\"LANCEDB_API_KEY\"],\n region=os.environ.get(\"LANCEDB_REGION\", \"us-east-1\"),\n host_override=os.environ[\"LANCEDB_HOST_OVERRIDE\"], # https://your-endpoint\n sql_host_override=os.environ[\"LANCEDB_SQL_HOST_OVERRIDE\"],\n)\n"; export const PyConnectObjectStorage = "import lancedb\n\nuri = \"s3://your-bucket/path\"\n# You can also use \"gs://your-bucket/path\" or \"az://your-container/path\".\ndb = lancedb.connect(uri)\n"; @@ -24,9 +24,9 @@ export const TsConnectEnterpriseQuickstart = "const setting = (name: string) => export const TsConnectObjectStorage = "const uri = \"s3://your-bucket/path\";\n// You can also use \"gs://your-bucket/path\" or \"az://your-container/path\".\nconst db = await lancedb.connect(uri);\n"; -export const TsNamespaceAdminOps = "const db = await lancedb.connectNamespace(\"dir\", { root: \"./local_lancedb\" });\nconst namespace = [\"prod\", \"search\"];\n\nawait db.createNamespace([\"prod\"]);\nawait db.createNamespace([\"prod\", \"search\"]);\n\nconst childNamespaces = (await db.listNamespaces([\"prod\"])).namespaces;\nconsole.log(`Child namespaces under ${JSON.stringify(namespace)}:`, childNamespaces);\n// Child namespaces under [\"prod\",\"search\"]: [ 'search' ]\n\nconst metadata = await db.describeNamespace([\"prod\", \"search\"]);\nconsole.log(`Metadata for namespace ${JSON.stringify(namespace)}:`, metadata);\n\nawait db.dropNamespace([\"prod\", \"search\"], { mode: \"skip\" });\nawait db.dropNamespace([\"prod\"], { mode: \"skip\" });\n"; +export const TsNamespaceAdminOps = "async function collectNames(items: AsyncIterable): Promise {\n const names: string[] = [];\n for await (const name of items) names.push(name);\n return names;\n}\n\nconst db = await lancedb.connectNamespace(\"dir\", { root: \"./local_lancedb\" });\nconst namespace = [\"prod\", \"search\"];\n\nawait db.createNamespace([\"prod\"]);\nawait db.createNamespace([\"prod\", \"search\"]);\n\nconst childNamespaces = await collectNames(db.listNamespaces([\"prod\"]));\nconsole.log(\n `Child namespaces under ${JSON.stringify(namespace)}:`,\n childNamespaces,\n);\n// Child namespaces under [\"prod\",\"search\"]: [ 'search' ]\n\nconst metadata = await db.describeNamespace([\"prod\", \"search\"]);\nconsole.log(`Metadata for namespace ${JSON.stringify(namespace)}:`, metadata);\n\nawait db.dropNamespace([\"prod\", \"search\"], { mode: \"skip\" });\nawait db.dropNamespace([\"prod\"], { mode: \"skip\" });\n"; -export const TsNamespaceTableOps = "const db = await lancedb.connectNamespace(\"dir\", { root: \"./local_lancedb\" });\n\n// Create namespace tree: prod/search and prod/recommendations\nawait db.createNamespace([\"prod\"], { mode: \"exist_ok\" });\nawait db.createNamespace([\"prod\", \"search\"], { mode: \"exist_ok\" });\nawait db.createNamespace([\"prod\", \"recommendations\"], { mode: \"exist_ok\" });\n\nawait db.createTable(\n \"user\",\n [{ id: 1, vector: [0.1, 0.2], name: \"alice\" }],\n [\"prod\", \"search\"],\n { mode: \"create\" }, // use \"overwrite\" only if you want to replace existing table\n);\n\nawait db.createTable(\n \"user\",\n [{ id: 2, vector: [0.3, 0.4], name: \"bob\" }],\n [\"prod\", \"recommendations\"],\n { mode: \"create\" },\n);\n\n// Verify\nconsole.log((await db.listNamespaces()).namespaces); // [\"prod\"]\nconsole.log((await db.listNamespaces([\"prod\"])).namespaces); // [\"recommendations\", \"search\"]\nconsole.log(await db.tableNames([\"prod\", \"search\"])); // [\"user\"]\nconsole.log(await db.tableNames([\"prod\", \"recommendations\"])); // [\"user\"]\n"; +export const TsNamespaceTableOps = "async function collectNames(items: AsyncIterable): Promise {\n const names: string[] = [];\n for await (const name of items) names.push(name);\n return names;\n}\n\nconst db = await lancedb.connectNamespace(\"dir\", { root: \"./local_lancedb\" });\n\n// Create namespace tree: prod/search and prod/recommendations\nawait db.createNamespace([\"prod\"], { mode: \"exist_ok\" });\nawait db.createNamespace([\"prod\", \"search\"], { mode: \"exist_ok\" });\nawait db.createNamespace([\"prod\", \"recommendations\"], { mode: \"exist_ok\" });\n\nawait db.createTable(\n \"user\",\n [{ id: 1, vector: [0.1, 0.2], name: \"alice\" }],\n [\"prod\", \"search\"],\n { mode: \"create\" }, // use \"overwrite\" only if you want to replace existing table\n);\n\nawait db.createTable(\n \"user\",\n [{ id: 2, vector: [0.3, 0.4], name: \"bob\" }],\n [\"prod\", \"recommendations\"],\n { mode: \"create\" },\n);\n\n// Verify\nconsole.log(await collectNames(db.listNamespaces())); // [\"prod\"]\nconsole.log(await collectNames(db.listNamespaces([\"prod\"]))); // [\"recommendations\", \"search\"]\nconsole.log(await db.tableNames([\"prod\", \"search\"])); // [\"user\"]\nconsole.log(await db.tableNames([\"prod\", \"recommendations\"])); // [\"user\"]\n"; export const RsConnect = "let db = connect(\"ex_lancedb\").execute().await?;\n"; diff --git a/docs/web/snippets/function.mdx b/docs/web/snippets/function.mdx index fb7f907c3..740dc397a 100644 --- a/docs/web/snippets/function.mdx +++ b/docs/web/snippets/function.mdx @@ -6,7 +6,7 @@ export const PyFunctionClassApply = "from lancedb import col\n\napplication = em export const PyFunctionClassDefine = "@udf(pip=[\"sentence-transformers\"])\nclass embed:\n def __init__(self, model: str = \"all-MiniLM-L6-v2\"):\n # Runs on the worker, once per instance: only the worker needs the package.\n from sentence_transformers import SentenceTransformer\n\n self.model = SentenceTransformer(model)\n\n def __call__(self, text: str) -> list[float]:\n return self.model.encode(text).tolist()\n\n def close(self):\n del self.model\n"; -export const PyFunctionConnect = "import os\n\nimport lancedb\n\ndb = lancedb.connect(\n uri=os.environ[\"LANCEDB_URI\"], # db://your-database\n api_key=os.environ[\"LANCEDB_API_KEY\"],\n region=os.environ.get(\"LANCEDB_REGION\", \"us-east-1\"),\n host_override=os.environ[\"LANCEDB_HOST_OVERRIDE\"],\n)\n"; +export const PyFunctionConnect = "import os\n\nimport lancedb\n\ndb = lancedb.connect(\n uri=os.environ[\"LANCEDB_URI\"], # db://your-database\n api_key=os.environ[\"LANCEDB_API_KEY\"],\n region=os.environ.get(\"LANCEDB_REGION\", \"us-east-1\"),\n host_override=os.environ[\"LANCEDB_HOST_OVERRIDE\"],\n sql_host_override=os.environ[\"LANCEDB_SQL_HOST_OVERRIDE\"],\n)\n"; export const PyFunctionCreateTable = "import pyarrow as pa\n\nschema = pa.schema([pa.field(\"value\", pa.float64(), nullable=False)])\ntable = db.create_table(\n \"function_demo\",\n [{\"value\": 1.0}, {\"value\": 2.0}, {\"value\": 3.0}],\n schema=schema,\n)\n"; diff --git a/docs/web/sql/ddl/function.mdx b/docs/web/sql/ddl/function.mdx index 1593c2d7b..854b66005 100644 --- a/docs/web/sql/ddl/function.mdx +++ b/docs/web/sql/ddl/function.mdx @@ -1,7 +1,7 @@ --- title: "Function" sidebarTitle: "Function" -description: "CREATE FUNCTION, SHOW FUNCTIONS, and DROP FUNCTION." +description: "Create, describe, alter, list, and drop registered functions with SQL." icon: function keywords: ["create function", "drop function", "show functions", "sql", "udf"] --- @@ -60,6 +60,35 @@ version is unaffected. the function is baked into an image before the version becomes usable. +## Artifact registration {#artifact} + +`CREATE FUNCTION name FROM ''` accepts the artifact, +signature, runtime, optional secret bindings, and result stability produced by +the SDK's registration request. Omit its `name` property because the SQL +identifier supplies it. Escape apostrophes in the JSON as doubled apostrophes. + +By default, the command waits for its internal image-build job to publish the +version, then returns one UTF-8 `definition` column containing FunctionVersion +JSON. `WITH (wait = false)` instead returns the existing image-build registry +job ID in `job_id`, with `operation` and `status` columns. The SDK's ordinary +registration handle has type `query` and uses the SQL query UUID as its single +`id`; the internal image-build job has type `create_function`. + +## DESCRIBE and ALTER FUNCTION {#describe-alter} + +```sql +DESCRIBE FUNCTION embed WITH (version = '1', include_source = true) +``` + +DESCRIBE returns one `definition` JSON column for the exact version. +`include_source` defaults to false; source-inclusive reads report an omission +when the stored artifact exceeds the inline source limit. + +`ALTER FUNCTION name FROM ''` accepts a `reference` with +name, namespace, object ID, location, version, and manifest digest, plus optional +`metadata` and `disabled` updates. The command checks that reference before +appending a revision and returns the resulting `definition` JSON. + ## Call it {#call} Once registered, a function is callable by name in any query, the same as a @@ -78,29 +107,27 @@ from it. ```sql SHOW FUNCTIONS +SHOW FUNCTIONS IN "analytics" WITH (name = 'embed', limit = 100) +SHOW FUNCTIONS WITH (after_name = 'embed', after_version = '1', limit = 100) ``` -Lists registered functions and their versions. Since a version is immutable -and a column pins the version that fills it, this is how you find which -versions are still in use before dropping one. +Rows are ordered lexicographically by name and version. `after_name` and +`after_version` must be supplied together for keyset pagination. `limit` bounds +the response; `include_definition = false` omits the optional definition JSON. ## DROP FUNCTION {#drop} ```sql DROP FUNCTION embed WITH (version = '1') DROP FUNCTION IF EXISTS embed WITH (version = '1') -DROP FUNCTION embed WITH (version = '1', cascade = 'true') ``` -`version` is required: a function is a set of immutable versions, and -dropping is always of one specific version, never "the function" as a whole. +`version` must identify an existing revision. DROP releases the function name +and its object binding; physical cleanup continues separately. The result +reports `dropped` or `not_found`. `IF EXISTS` allows a missing name. - -Without `cascade = 'true'`, dropping a version that a column still pins fails -rather than breaking that column silently. Passing it drops the version and -lets columns pinned to it stop refreshing. They keep their last-filled -values, per [Guarantees](/basics/guarantees#jobs). - +All function object commands retain shared results even when general query-result +copying is disabled, so asynchronous polling can resume on another query node. ## What to read next {#next} diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index d50927936..c43e71744 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -975,7 +975,8 @@ class DBConnection(EnforceOverrides): """Submit a scalar Python UDF for building and registration. The server-side job builds the OCI image, then registers the completed - artifact. Waiting on the job returns the immutable Function version. + artifact. The returned job has type ``query`` and the SQL query UUID as + its ID. Waiting on it returns the immutable Function version. LanceDB Embedded connections raise ``NotImplementedError``. """ raise NotImplementedError( @@ -1033,10 +1034,10 @@ class DBConnection(EnforceOverrides): """Drop a Function name and the object it was bound to. The requested version must exist in the currently named object. Returns - True when the name was removed and False when it was absent. The - object's content is deleted with it, which may finish after this - returns; use :meth:`drop_function_async` to wait for that. Local - connections raise NotImplementedError. + True when the catalog binding was removed and False when it was absent. + Physical cleanup continues in the background; use + :meth:`drop_function_async` to access the query job. Local connections + raise NotImplementedError. """ raise NotImplementedError( "Function catalog operations are not supported for this connection type" @@ -1049,13 +1050,12 @@ class DBConnection(EnforceOverrides): version: str, namespace_path: Optional[List[str]] = None, ) -> Tuple[bool, Job]: - """Drop a Function name and return its cleanup job. + """Drop a Function binding and return its SQL query job. - The name is unbound before this returns; the object's content may still - be being deleted. Call :meth:`Job.wait` to wait for that to finish. When - the server deletes inline, or when nothing was bound, the returned job - is already finished and has no id. LanceDB Embedded connections raise - NotImplementedError. + The boolean reports catalog removal. The job has type ``query`` and its + ID is the SQL query UUID. Repeated waits confirm that completion; + physical cleanup continues in the background. LanceDB Embedded + connections raise NotImplementedError. """ raise NotImplementedError( "Function catalog operations are not supported for this connection type" @@ -3277,7 +3277,8 @@ class AsyncConnection(object): """Submit a scalar Python UDF for building and registration. The server-side job builds the OCI image, then registers the completed - artifact. Waiting on the job returns the immutable Function version. + artifact. The returned job has type ``query`` and the SQL query UUID as + its ID. Waiting on it returns the immutable Function version. ``secrets`` is a sequence of [EnvVarSecret][lancedb.secrets.EnvVarSecret], each naming a Secret and the environment variable its value arrives in. ``namespace_path`` is @@ -3347,11 +3348,11 @@ class AsyncConnection(object): version: str, namespace_path: Optional[List[str]] = None, ) -> Tuple[bool, AsyncJob]: - """Drop a Function name and return its cleanup job. + """Drop a Function binding and return its SQL query job. - The name is unbound before this returns; the object's content may still - be being deleted. Await :meth:`AsyncJob.wait` to wait for that to - finish. + The boolean reports catalog removal. The job has type ``query`` and its + ID is the SQL query UUID. Repeated waits confirm that completion; + physical cleanup continues in the background. """ dropped, job = await self._inner.drop_function_async( name, version, namespace_path diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index e1e3784d6..6240c362a 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -245,36 +245,6 @@ def test_bindings_may_not_collide_with_plain_configuration(): ) -def test_a_binding_envelope_reaches_the_service_for_it_to_judge(): - """Binding rules are the service's: it owns the runtime the names land in. - - The client sends what it was given, so a rule it duplicated could disagree - with the service's without either side noticing. What is checked here is - that the envelope arrives intact -- the shape the service judges is the - shape the caller wrote. - """ - with _mock_remote_function_catalog() as (host, state): - db = lancedb.connect( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - bindings = [ - EnvVarSecret(secret_name=f"secret-{index}", env_variable=f"TOKEN_{index}") - for index in range(17) - ] - db.create_function(normalize_score, secrets=bindings) - - sent = state["requests"][0][1] - assert len(sent["secret_bindings"]) == 17 - assert { - "kind": "env", - "variable": "TOKEN_0", - "secret_ref": {"name": "secret-0"}, - } in sent["secret_bindings"] - - _SECRET_DEBUG_LOG_SOURCE = """ import http.server import json @@ -1846,80 +1816,6 @@ def _mock_remote_function_catalog(): thread.join() -def test_remote_registration_job_and_exact_version_reopen_round_trip(): - with _mock_remote_function_catalog() as (host, state): - db = lancedb.connect( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - registration = db.create_function_async(normalize_score) - assert registration.id == "job-register" - assert registration.job_type is None - assert registration.status() == "finished" - assert registration.job_type == "create_function" - created = registration.wait() - assert registration.job_type == "create_function" - reopened = db.get_function("normalize_score", version=created.version) - - assert created == reopened - assert reopened.name == "normalize_score" - assert reopened.version == FUNCTION_VERSION - assert ( - reopened.source_digest == normalize_score.registration_request.artifact.digest - ) - create_request = state["requests"][0][1] - expected_request = json.loads( - normalize_score.registration_request.to_canonical_json() - ) - expected_request.pop("name") - assert create_request == expected_request - - -def test_remote_registration_sends_bindings_and_never_a_credential(): - with _mock_remote_function_catalog() as (host, state): - db = lancedb.connect( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - created = db.create_function( - analyze_caption, - secrets=[ - EnvVarSecret(secret_name="openai-prod", env_variable="OPENAI_API_KEY") - ], - ) - - assert list(created.secret_bindings) == [ - SecretBinding( - kind="env", - variable="OPENAI_API_KEY", - secret_ref=SecretReference(name="openai-prod"), - ) - ] - path, create_request = state["requests"][0] - assert path == "/v1/function/analyze_caption/create" - assert create_request["secret_bindings"] == [ - { - "kind": "env", - "variable": "OPENAI_API_KEY", - "secret_ref": {"name": "openai-prod"}, - } - ] - # The request names a Secret and carries nothing that could be one. The - # Function's own name is the path identifier rather than a body field, so - # it is the one key the body does not repeat. - expected = json.loads( - analyze_caption.bind_secrets( - [EnvVarSecret(secret_name="openai-prod", env_variable="OPENAI_API_KEY")] - ).to_canonical_json() - ) - assert expected.pop("name") == "analyze_caption" - assert create_request == expected - - def test_remote_secret_verbs_round_trip(): with _mock_remote_function_catalog() as (host, state): db = lancedb.connect( @@ -1963,7 +1859,7 @@ def test_building_a_binding_contacts_no_server(): assert state["requests"] == [] -def test_registered_vectorized_blob_udf_hydrates_executes_and_publishes(tmp_path): +def test_vectorized_blob_udf_executes_packaged_artifact_and_publishes(tmp_path): @udf( input_schema=pa.schema([lancedb.blob("image", nullable=True)]), output_schema=lancedb.blob("copy", nullable=False), @@ -1978,19 +1874,7 @@ def test_registered_vectorized_blob_udf_hydrates_executes_and_publishes(tmp_path type=pa.large_binary(), ) - with _mock_remote_function_catalog() as (host, state): - remote = lancedb.connect( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - remote.create_function(copy_blobs) - registration = state["requests"][0][1] - - definition = json.loads(copy_blobs.registration_request.to_canonical_json()) - assert registration["artifact"] == definition["artifact"] - assert registration["signature"] == definition["signature"] + registration = json.loads(copy_blobs.registration_request.to_canonical_json()) local = lancedb.connect(tmp_path) source = local.create_table( @@ -2049,196 +1933,6 @@ def test_registered_vectorized_blob_udf_hydrates_executes_and_publishes(tmp_path ] -def test_blocking_remote_registration_returns_function_version(): - with _mock_remote_function_catalog() as (host, state): - db = lancedb.connect( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - created = db.create_function(normalize_score) - - assert created.name == "normalize_score" - assert created.version == FUNCTION_VERSION - assert created.source_digest == normalize_score.registration_request.artifact.digest - assert [path for path, _ in state["requests"]] == [ - "/v1/function/normalize_score/create", - "/v1/jobs/describe", - ] - - -def test_remote_list_functions_paginates_and_returns_typed_versions(): - with _mock_remote_function_catalog() as (host, state): - db = lancedb.connect( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - created = db.create_function(normalize_score) - state["requests"].clear() - functions = list(db.list_functions()) - - assert functions == [created] - assert functions[0].source_digest == ( - normalize_score.registration_request.artifact.digest - ) - assert state["requests"] == [ - ("/v1/namespace/$/function/list", {"include_definition": "true"}), - ( - "/v1/namespace/$/function/list", - {"include_definition": "true", "page_token": "next"}, - ), - ] - - -@pytest.mark.asyncio -async def test_async_remote_list_functions_returns_typed_versions(): - with _mock_remote_function_catalog() as (host, state): - db = await lancedb.connect_async( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - registration = await db.create_function_async(normalize_score) - created = await registration.wait() - state["requests"].clear() - functions = [item async for item in db.list_functions()] - - assert functions == [created] - assert functions[0].source_digest == ( - normalize_score.registration_request.artifact.digest - ) - assert [path for path, _ in state["requests"]] == [ - "/v1/namespace/$/function/list", - "/v1/namespace/$/function/list", - ] - - -def test_remote_drop_function_sends_exact_version(): - with _mock_remote_function_catalog() as (host, state): - db = lancedb.connect( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - assert db.drop_function("normalize_score", version=FUNCTION_VERSION) is True - - assert state["requests"] == [ - ( - "/v1/function/normalize_score/drop", - {"version": FUNCTION_VERSION}, - ) - ] - - -@pytest.mark.asyncio -async def test_async_remote_drop_function_sends_exact_version(): - with _mock_remote_function_catalog() as (host, state): - db = await lancedb.connect_async( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - assert ( - await db.drop_function("normalize_score", version=FUNCTION_VERSION) is True - ) - - assert state["requests"] == [ - ( - "/v1/function/normalize_score/drop", - {"version": FUNCTION_VERSION}, - ) - ] - - -def test_remote_function_verbs_address_a_namespace_in_the_path(): - """A Function's namespace is part of the identifier each route addresses; - no body carries it, so a namespaced request differs only in its path.""" - namespace = ["analytics", "features"] - with _mock_remote_function_catalog() as (host, state): - db = lancedb.connect( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - created = db.create_function(normalize_score, namespace_path=namespace) - # The returned parts are what the other verbs take, so a created - # Function reopens from its own fields. - assert created.name == "normalize_score" - assert created.namespace_path == tuple(namespace) - reopened = db.get_function( - created.name, - version=created.version, - namespace_path=list(created.namespace_path), - ) - assert reopened == created - assert reopened(value=lancedb.col("score")).function.namespace_path == tuple( - namespace - ) - list(db.list_functions(namespace_path=namespace)) - assert ( - db.drop_function( - "normalize_score", version=FUNCTION_VERSION, namespace_path=namespace - ) - is True - ) - dropped, _ = db.drop_function_async( - "normalize_score", version=FUNCTION_VERSION, namespace_path=namespace - ) - assert dropped is True - - assert [path for path, _ in state["requests"]] == [ - "/v1/function/analytics$features$normalize_score/create", - "/v1/jobs/describe", - "/v1/function/analytics$features$normalize_score/describe", - "/v1/namespace/analytics$features/function/list", - "/v1/namespace/analytics$features/function/list", - "/v1/function/analytics$features$normalize_score/drop", - "/v1/function/analytics$features$normalize_score/drop", - ] - - -@pytest.mark.asyncio -async def test_async_remote_function_verbs_address_a_namespace_in_the_path(): - namespace = ["analytics", "features"] - with _mock_remote_function_catalog() as (host, state): - db = await lancedb.connect_async( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - registration = await db.create_function_async( - normalize_score, namespace_path=namespace - ) - await registration.wait() - await db.get_function( - "normalize_score", version=FUNCTION_VERSION, namespace_path=namespace - ) - [item async for item in db.list_functions(namespace_path=namespace)] - assert ( - await db.drop_function( - "normalize_score", version=FUNCTION_VERSION, namespace_path=namespace - ) - is True - ) - - assert [path for path, _ in state["requests"]] == [ - "/v1/function/analytics$features$normalize_score/create", - "/v1/jobs/describe", - "/v1/function/analytics$features$normalize_score/describe", - "/v1/namespace/analytics$features/function/list", - "/v1/namespace/analytics$features/function/list", - "/v1/function/analytics$features$normalize_score/drop", - ] - - def test_a_function_namespace_path_must_be_a_list_of_segments(): """A bare string is refused rather than read as a sequence of one-character segments, and nothing reaches the service.""" diff --git a/python/python/tests/test_listing.py b/python/python/tests/test_listing.py index 4cf37b631..7b916dadc 100644 --- a/python/python/tests/test_listing.py +++ b/python/python/tests/test_listing.py @@ -84,7 +84,7 @@ async def advance(items, asynchronous): return await items.__anext__() if asynchronous else next(items) -KINDS = ["secrets", "views", "jobs", "functions"] +KINDS = ["secrets", "views", "jobs"] @pytest.mark.asyncio diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 1b1920414..c7559f999 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -646,9 +646,10 @@ impl Connection { /// in `namespace_path`, which is empty for the root namespace. /// /// The server-side job builds the OCI image, then registers the completed - /// artifact. Waiting on the returned typed job yields the durable - /// [`crate::function::FunctionVersion`]. Creation is supported only in LanceDB Distributed. - /// LanceDB Embedded databases return [`Error::NotSupported`]. + /// artifact. The returned job has type `query` and the SQL query UUID as its + /// ID. Waiting on the returned typed job yields the durable + /// [`crate::function::FunctionVersion`]. Creation is remote-only. + /// Local databases return [`Error::NotSupported`]. /// pub async fn create_function_async( &self, @@ -730,12 +731,11 @@ impl Connection { .await } - /// Start dropping a Function and return its cleanup job. + /// Drop a Function binding and return its SQL query handle. /// - /// The name is unbound before this returns; the object's content may still be being - /// deleted. Await [`Job::wait`][crate::job::Job::wait] to wait for that to finish. When - /// the server deletes inline, or when nothing was bound, the returned job is already - /// finished and has no id. LanceDB Embedded databases return [`Error::NotSupported`]. + /// The boolean reports catalog removal. The query job retains its query ID, + /// and repeated waits confirm completion. Physical cleanup continues in the + /// background. Local databases return [`Error::NotSupported`]. pub async fn drop_function_async( &self, name: impl AsRef, diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index 4528257d1..f5fba221d 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -460,10 +460,9 @@ pub trait Database: ) -> Result { function_catalog_not_supported() } - /// Start dropping a Function and return a handle to the cleanup job. - /// - /// Backends without asynchronous cleanup complete the drop before returning an - /// already-finished job. + /// Drop a Function binding and return its completion handle. + /// SQL backends return a query job confirming catalog removal; physical + /// cleanup continues separately. async fn drop_function_async( &self, name: &str, diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 445961205..b7a641dd3 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -25,16 +25,12 @@ use crate::database::{ ExecuteQueryRequest, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest, }; use crate::error::Result; -use crate::function::{ - FunctionArtifactRequest, FunctionRegistrationRequest, FunctionSignature, FunctionVersion, - PythonRuntimeSpec, -}; +use crate::function::{FunctionRegistrationRequest, FunctionVersion}; use crate::job::Job; use crate::listing::{ListingOptions, ListingPage}; use crate::materialized_view::CreateMaterializedViewRequest; use crate::remote::job::{PauseJobResponse, RemoteJob, ResumeJobResponse, job_state_to_client}; use crate::remote::util::stream_as_body; -use crate::secrets::SecretBinding; use crate::secrets::SecretInfo; use crate::table::BaseTable; use crate::utils::{ @@ -475,6 +471,12 @@ impl RemoteDatabase { } } + fn sql_client(&self) -> Result<&SqlClient> { + self.sql_client.as_ref().ok_or_else(|| Error::NotSupported { + message: "SQL is unavailable for this remote database client".into(), + }) + } + /// Post a request whose body carries a credential. /// /// Shared by the create and alter verbs, which declare their own request @@ -668,14 +670,6 @@ fn build_namespace_identifier(namespace: &[String]) -> Result { } } -/// Build a full path to a function. Note that we do not need to percent-encode -/// this when using it in a URL because the validation functions ensure that -/// only URL-safe characters are present. -fn build_function_identifier(name: &str, namespace: &[String]) -> Result { - validate_function_name(name)?; - build_compound_identifier(name, namespace) -} - /// Build a full path to a secret. Note that we do not need to percent-encode /// this when using it in a URL because the validation functions ensure that /// only URL-safe characters are present. @@ -736,49 +730,6 @@ struct RemoteListJobsResponse { page_token: Option, } -#[derive(serde::Deserialize)] -struct RemoteListedFunctionVersion { - definition: FunctionVersion, -} - -#[derive(serde::Deserialize)] -struct RemoteListFunctionsResponse { - #[serde(default)] - functions: Vec, - #[serde(default)] - page_token: Option, -} - -#[derive(serde::Deserialize)] -struct RemoteDropFunctionResponse { - dropped: bool, -} - -/// The create body: every field of [`FunctionRegistrationRequest`] except the -/// name, which is the path identifier. -/// -/// A struct rather than a literal listing the fields, so that the compiler -/// decides what reaches the service. A field the registration request grows is -/// a build error here until it is handled; a literal would simply not send it. -#[derive(serde::Serialize)] -struct RemoteCreateFunctionRequest<'a> { - artifact: &'a FunctionArtifactRequest, - signature: &'a FunctionSignature, - runtime: &'a PythonRuntimeSpec, - /// Absent when the Function binds nothing, so such a client sends what a - /// client without bindings sends. - /// - /// A service that does not know the field ignores it: registration - /// succeeds, the returned version carries no bindings, and the Function - /// fails at execution with the variable unset. [`ServerVersion`] is how - /// this codebase refuses a feature the service is too old for; it is held - /// per table, so gating a database-level call is follow-up work. - /// - /// [`ServerVersion`]: super::db::ServerVersion - #[serde(skip_serializing_if = "<[SecretBinding]>::is_empty")] - secret_bindings: &'a [SecretBinding], -} - /// Create a Secret under a name the database does not yet hold. /// /// Declared separately from the alter request although the two are identical @@ -976,29 +927,11 @@ impl Database for RemoteDatabase { request: FunctionRegistrationRequest, namespace_path: &[String], ) -> Result> { - let function_id = build_function_identifier(&request.name, namespace_path)?; - let req = self - .client - .post(&format!("/v1/function/{function_id}/create")) - .json(&RemoteCreateFunctionRequest { - artifact: &request.artifact, - signature: &request.signature, - runtime: &request.runtime, - secret_bindings: &request.secret_bindings, - }); - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let body = response.text().await.err_to_http(request_id.clone())?; - let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { - source: "Function registration response did not contain a valid job_id".into(), - request_id, - status_code: Some(status), - })?; - Ok(Job::new_typed(Box::new(RemoteJob::new( - self.client.clone(), - job_id, - )))) + validate_function_name(&request.name)?; + validate_namespace(namespace_path)?; + self.sql_client()? + .register_function(request, namespace_path) + .await } async fn get_function( @@ -1007,46 +940,22 @@ impl Database for RemoteDatabase { version: &str, namespace_path: &[String], ) -> Result { - let function_id = build_function_identifier(name, namespace_path)?; - let req = self - .client - .post(&format!("/v1/function/{function_id}/describe")) - .json(&serde_json::json!({ - "version": version, - })); - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - response.json().await.err_to_http(request_id) + validate_function_name(name)?; + validate_namespace(namespace_path)?; + self.sql_client()? + .describe_function(name, version, namespace_path) + .await } async fn list_functions( &self, namespace_path: &[String], options: ListingOptions, - ) -> Result> { - let namespace_id = build_namespace_identifier(namespace_path)?; - let path = format!("/v1/namespace/{namespace_id}/function/list"); - let mut req = self.client.get(&path); - if let Some(token) = options.page_token { - req = req.query(&[("page_token", token)]); - } - if let Some(limit) = options.page_limit { - req = req.query(&[("limit", limit)]); - } - req = req.query(&[("include_definition", true)]); - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - let response: RemoteListFunctionsResponse = - response.json().await.err_to_http(request_id)?; - let items = response - .functions - .into_iter() - .map(|listed| listed.definition) - .collect(); - Ok(ListingPage { - items, - page_token: response.page_token, - }) + ) -> Result> { + validate_namespace(namespace_path)?; + self.sql_client()? + .list_functions(namespace_path, options) + .await } async fn drop_function( @@ -1055,10 +964,14 @@ impl Database for RemoteDatabase { version: &str, namespace_path: &[String], ) -> Result { - Ok(self - .drop_function_async(name, version, namespace_path) - .await? - .0) + validate_function_name(name)?; + validate_namespace(namespace_path)?; + let (dropped, job) = self + .sql_client()? + .drop_function(name, version, namespace_path) + .await?; + job.wait().await?; + Ok(dropped) } async fn drop_function_async( @@ -1067,43 +980,11 @@ impl Database for RemoteDatabase { version: &str, namespace_path: &[String], ) -> Result<(bool, Job)> { - let function_id = build_function_identifier(name, namespace_path)?; - let req = self - .client - .post(&format!("/v1/function/{function_id}/drop")) - .json(&serde_json::json!({ - "version": version, - })); - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let body = response.text().await.err_to_http(request_id.clone())?; - let dropped: RemoteDropFunctionResponse = - serde_json::from_str(&body).map_err(|source| Error::Http { - source: Box::new(source), - request_id: request_id.clone(), - status_code: Some(status), - })?; - let job = match status { - StatusCode::OK => Job::new_done(), - StatusCode::ACCEPTED => { - let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { - source: "asynchronous Function drop response did not contain a valid job_id" - .into(), - request_id, - status_code: Some(status), - })?; - Job::new(Box::new(RemoteJob::new(self.client.clone(), job_id))) - } - _ => { - return Err(Error::Http { - source: "Function drop must return 200 OK or 202 Accepted".into(), - request_id, - status_code: Some(status), - }); - } - }; - Ok((dropped.dropped, job)) + validate_function_name(name)?; + validate_namespace(namespace_path)?; + self.sql_client()? + .drop_function(name, version, namespace_path) + .await } async fn create_storage_integration( @@ -2142,57 +2023,6 @@ mod tests { } } - #[tokio::test] - async fn test_drop_function_async_returns_job() { - let db = super::RemoteDatabase::new_mock(|request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/function/embed/drop"); - http::Response::builder() - .status(202) - .body(serde_json::json!({"dropped": true, "job_id": "j1-fn-drop"}).to_string()) - .unwrap() - }); - let (dropped, job) = db.drop_function_async("embed", "1", &[]).await.unwrap(); - assert!(dropped); - assert_eq!(job.id(), Some("j1-fn-drop")); - } - - /// An unbound name and an inline deletion both answer `200`: the name is gone and - /// nothing is left to wait for, so the job is already finished. - #[tokio::test] - async fn test_drop_function_async_completed_inline() { - let db = super::RemoteDatabase::new_mock(|_| { - http::Response::builder() - .status(200) - .body(serde_json::json!({"dropped": false}).to_string()) - .unwrap() - }); - let (dropped, job) = db.drop_function_async("embed", "1", &[]).await.unwrap(); - assert!(!dropped); - assert_eq!(job.id(), None); - assert_eq!(job.status().await.unwrap(), "finished"); - job.wait().await.unwrap(); - } - - #[tokio::test] - async fn test_drop_function_async_rejects_incomplete_acceptance() { - for body in [ - r#"{"dropped":true}"#, - r#"{"dropped":true,"job_id":""}"#, - r#"{"dropped":true,"job_id":null}"#, - ] { - let db = super::RemoteDatabase::new_mock(move |_| { - http::Response::builder().status(202).body(body).unwrap() - }); - let error = db - .drop_function_async("embed", "1", &[]) - .await - .err() - .unwrap(); - assert!(error.to_string().contains("valid job_id")); - } - } - #[tokio::test] async fn test_retries() { // We'll record the request_id here, to check it matches the one in the error. @@ -3874,7 +3704,7 @@ mod tests { assert!(items.try_next().await.unwrap().is_none()); assert!(items.try_next().await.unwrap().is_none()); } - for resource in ["secret", "view", "function", "jobs"] { + for resource in ["secret", "view", "jobs"] { let count = Arc::new(AtomicUsize::new(0)); let conn = Connection::new_with_handler({ let count = count.clone(); @@ -3897,9 +3727,6 @@ mod tests { let query = request.url().query_pairs().collect::>(); assert_eq!(query.get("page_token").unwrap(), token); assert_eq!(query.get("limit").unwrap(), "2"); - if resource == "function" { - assert_eq!(query.get("include_definition").unwrap(), "true"); - } } let (key, item) = match resource { "secret" => ("secrets", serde_json::json!({"name": "a"})), @@ -3907,10 +3734,6 @@ mod tests { "jobs", serde_json::json!({"job_id": "a", "table": "t", "job_type": "create_index", "state": "in_progress", "created_at_millis": 1}), ), - "function" => ( - "functions", - serde_json::json!({"name": "embed", "version": "1", "definition": serde_json::from_str::(include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json")).unwrap()}), - ), _ => ("views", serde_json::json!("a")), }; let items = if page == 1 { @@ -3928,7 +3751,6 @@ mod tests { match resource { "secret" => check(conn.list_secrets(&namespace, options), count).await, "view" => check(conn.list_views(&namespace, options), count).await, - "function" => check(conn.list_functions(&namespace, options), count).await, _ => check(conn.list_jobs(options), count).await, } } @@ -4389,27 +4211,6 @@ mod tests { )); } - /// Reject relative segments and their encoded forms before sending a request, - /// so a function name cannot change the route through URL normalization. - #[tokio::test] - async fn test_a_relative_segment_function_name_is_refused() { - use std::sync::{Arc, Mutex}; - for name in [".", "..", "%2E%2E"] { - let reached = Arc::new(Mutex::new(false)); - let flag = reached.clone(); - let conn = Connection::new_with_handler(move |_| { - *flag.lock().unwrap() = true; - http::Response::builder().status(200).body("{}").unwrap() - }); - let error = conn - .drop_function(name, "fv_1", &[]) - .await - .expect_err("a dot-only Function name must be refused"); - assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); - assert!(!*reached.lock().unwrap(), "{name:?} reached the transport"); - } - } - /// Only `.` and `..` are relative segments. `...` and longer runs are /// ordinary and address perfectly well, so refusing them would make an /// object that works today stop working on upgrade. This pins that. @@ -4792,281 +4593,6 @@ mod tests { } } - #[tokio::test] - async fn test_create_function_async_sends_canonical_request_and_decodes_typed_job() { - const REQUEST: &str = include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json" - ); - const FUNCTION_JOB: &str = - include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json"); - let mut expected: serde_json::Value = serde_json::from_str(REQUEST).unwrap(); - expected.as_object_mut().unwrap().remove("name"); - let conn = Connection::new_with_handler(move |request| match request.url().path() { - "/v1/function/normalize_score/create" => { - assert_eq!(request.method(), &reqwest::Method::POST); - let body: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - assert_eq!(body, expected); - http::Response::builder() - .status(202) - .body(r#"{"job_id":"job-function-1"}"#) - .unwrap() - } - "/v1/jobs/describe" => http::Response::builder() - .status(200) - .body(FUNCTION_JOB) - .unwrap(), - path => panic!("unexpected path: {path}"), - }); - let request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap(); - let job = conn.create_function_async(request, &[]).await.unwrap(); - assert_eq!(job.id(), Some("job-function-1")); - let version = job.wait().await.unwrap(); - assert_eq!(version.name(), "embed"); - assert_eq!(version.version(), "1"); - } - - #[tokio::test] - async fn test_get_function_requires_and_sends_exact_version() { - const VERSION: &str = include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json" - ); - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), "/v1/function/embed/describe"); - let body: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - assert_eq!(body, serde_json::json!({"version": "1"})); - http::Response::builder().status(200).body(VERSION).unwrap() - }); - let version = conn.get_function("embed", "1", &[]).await.unwrap(); - assert_eq!(version.name(), "embed"); - assert_eq!(version.version(), "1"); - } - - #[tokio::test] - async fn test_list_functions_requests_definitions_and_paginates() { - const VERSION: &str = include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json" - ); - let version: serde_json::Value = serde_json::from_str(VERSION).unwrap(); - let page = Arc::new(AtomicUsize::new(0)); - let conn = Connection::new_with_handler(move |request| { - assert_eq!(request.method(), &reqwest::Method::GET); - assert_eq!(request.url().path(), "/v1/namespace/$/function/list"); - let query = request.url().query_pairs().collect::>(); - assert_eq!(query.get("include_definition").unwrap(), "true"); - match page.fetch_add(1, Ordering::SeqCst) { - 0 => { - assert!(!query.contains_key("page_token")); - http::Response::builder() - .status(200) - .body(r#"{"functions": [], "page_token": "next"}"#.to_string()) - .unwrap() - } - _ => { - assert_eq!(query.get("page_token").unwrap(), "next"); - http::Response::builder() - .status(200) - .body( - serde_json::json!({ - "functions": [{ - "name": "embed", - "version": "1", - "definition": version.clone(), - }], - }) - .to_string(), - ) - .unwrap() - } - } - }); - let functions = conn - .list_functions(&[], Default::default()) - .try_collect::>() - .await - .unwrap(); - assert_eq!(functions.len(), 1); - assert_eq!(functions[0].name(), "embed"); - assert_eq!(functions[0].version(), "1"); - } - - #[tokio::test] - async fn test_list_functions_stops_on_an_empty_page_token() { - let requests = Arc::new(AtomicUsize::new(0)); - let seen = requests.clone(); - let conn = Connection::new_with_handler(move |request| { - seen.fetch_add(1, Ordering::SeqCst); - assert_eq!(request.method(), &reqwest::Method::GET); - assert_eq!(request.url().path(), "/v1/namespace/$/function/list"); - let query = request.url().query_pairs().collect::>(); - assert_eq!(query.get("include_definition").unwrap(), "true"); - assert!(!query.contains_key("page_token")); - http::Response::builder() - .status(200) - .body(r#"{"functions": [], "page_token": ""}"#) - .unwrap() - }); - - let functions = conn - .list_functions(&[], Default::default()) - .try_collect::>() - .await - .unwrap(); - assert!(functions.is_empty()); - assert_eq!(requests.load(Ordering::SeqCst), 1); - } - - #[tokio::test] - async fn test_list_functions_rejects_a_page_token_cycle() { - let page = Arc::new(AtomicUsize::new(0)); - let requests = page.clone(); - let conn = Connection::new_with_handler(move |request| { - assert_eq!(request.method(), &reqwest::Method::GET); - assert_eq!(request.url().path(), "/v1/namespace/$/function/list"); - let query = request.url().query_pairs().collect::>(); - assert_eq!(query.get("include_definition").unwrap(), "true"); - let next_page_token = match page.fetch_add(1, Ordering::SeqCst) { - 0 => { - assert!(!query.contains_key("page_token")); - "one" - } - 1 => { - assert_eq!(query.get("page_token").unwrap(), "one"); - "two" - } - 2 => { - assert_eq!(query.get("page_token").unwrap(), "two"); - "one" - } - page => panic!("unexpected page: {page}"), - }; - http::Response::builder() - .status(200) - .body( - serde_json::json!({ - "functions": [], - "page_token": next_page_token, - }) - .to_string(), - ) - .unwrap() - }); - - let mut items = conn.list_functions(&[], Default::default()); - let error = (&mut items).try_collect::>().await.unwrap_err(); - assert!( - matches!(&error, Error::Runtime { message } if message.contains("repeated a page_token")), - "got {error:?}" - ); - assert!(items.page_token().is_some()); - assert!(items.try_next().await.unwrap().is_none()); - assert_eq!(requests.load(Ordering::SeqCst), 3); - } - - #[tokio::test] - async fn test_drop_function_sends_exact_version_and_decodes_replay() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), "/v1/function/embed/drop"); - let body: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - assert_eq!(body, serde_json::json!({"version": "1"})); - http::Response::builder() - .status(200) - .body(r#"{"dropped":false}"#) - .unwrap() - }); - assert!(!conn.drop_function("embed", "1", &[]).await.unwrap()); - } - - /// A Function's namespace is addressed in the path the way a Secret's is: - /// every route takes the joined identifier and no body carries the - /// namespace, so a namespaced request differs from a root one only in its - /// path. - #[tokio::test] - async fn test_a_function_namespace_path_is_addressed_in_the_path() { - const REQUEST: &str = include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json" - ); - const VERSION: &str = include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json" - ); - let namespace = ["analytics".to_string(), "features".to_string()]; - let mut expected_create: serde_json::Value = serde_json::from_str(REQUEST).unwrap(); - expected_create.as_object_mut().unwrap().remove("name"); - let conn = Connection::new_with_handler(move |request| { - let body = request - .body() - .and_then(|body| body.as_bytes()) - .map(|bytes| serde_json::from_slice::(bytes).unwrap()); - match request.url().path() { - "/v1/function/analytics$features$normalize_score/create" => { - assert_eq!(body.unwrap(), expected_create); - http::Response::builder() - .status(202) - .body(r#"{"job_id":"job-function-1"}"#.to_string()) - .unwrap() - } - "/v1/function/analytics$features$embed/describe" => { - assert_eq!(body.unwrap(), serde_json::json!({"version": "1"})); - http::Response::builder() - .status(200) - .body(VERSION.to_string()) - .unwrap() - } - "/v1/function/analytics$features$embed/drop" => { - assert_eq!(body.unwrap(), serde_json::json!({"version": "1"})); - http::Response::builder() - .status(200) - .body(r#"{"dropped":true}"#.to_string()) - .unwrap() - } - // Listing is namespace-scoped, so the namespace is the whole - // identifier. - "/v1/namespace/analytics$features/function/list" => http::Response::builder() - .status(200) - .body(r#"{"functions":[]}"#.to_string()) - .unwrap(), - path => panic!("unexpected path: {path}"), - } - }); - let request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap(); - let job = conn - .create_function_async(request, &namespace) - .await - .unwrap(); - assert_eq!(job.id(), Some("job-function-1")); - conn.get_function("embed", "1", &namespace).await.unwrap(); - assert!( - conn.list_functions(&namespace, Default::default()) - .try_collect::>() - .await - .unwrap() - .is_empty() - ); - assert!(conn.drop_function("embed", "1", &namespace).await.unwrap()); - } - - /// Each namespace segment is checked before a request is built, so an - /// empty segment cannot collapse the identifier onto the parent namespace. - #[tokio::test] - async fn test_an_unaddressable_function_namespace_segment_is_refused() { - let conn = Connection::new_with_handler(|request| -> http::Response { - panic!("reached the transport: {}", request.url().path()) - }); - let namespace = ["analytics".to_string(), String::new()]; - assert!(conn.get_function("embed", "1", &namespace).await.is_err()); - assert!( - conn.list_functions(&namespace, Default::default()) - .try_collect::>() - .await - .is_err() - ); - assert!(conn.drop_function("embed", "1", &namespace).await.is_err()); - } - #[tokio::test] async fn test_conn_job_waits_to_done() { let polls = Arc::new(AtomicUsize::new(0)); diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index 880f61045..69cba2ef5 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -35,7 +35,9 @@ use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; +mod function; pub(super) mod materialized_view; +mod objects; mod stream; use stream::{StreamingQuery, StreamingQueryHandle, Submission}; diff --git a/rust/lancedb/src/remote/sql/function.rs b/rust/lancedb/src/remote/sql/function.rs new file mode 100644 index 000000000..7e468f14b --- /dev/null +++ b/rust/lancedb/src/remote/sql/function.rs @@ -0,0 +1,192 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::{ + SqlClient, + materialized_view::quote_identifier, + objects::{definition, literal, strings}, +}; +use crate::{ + database::ExecuteQueryRequest, + error::{Error, Result}, + function::{FunctionRegistrationRequest, FunctionVersion}, + job::Job, + listing::{ListingOptions, ListingPage}, + utils::{validate_function_name, validate_namespace}, +}; +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; + +impl SqlClient { + pub(in crate::remote) async fn register_function( + &self, + request: FunctionRegistrationRequest, + namespace: &[String], + ) -> Result> { + validate_function_name(&request.name)?; + validate_namespace(namespace)?; + let name = quote_identifier(&request.name); + let mut body = serde_json::to_value(request).map_err(|error| Error::Runtime { + message: error.to_string(), + })?; + body.as_object_mut().unwrap().remove("name"); + self.object_job( + ExecuteQueryRequest { + query: format!("CREATE FUNCTION {name} FROM {}", literal(&body.to_string())), + default_namespace_path: namespace.to_vec(), + parameters: None, + }, + definition, + ) + .await + } + pub(in crate::remote) async fn describe_function( + &self, + name: &str, + version: &str, + namespace: &[String], + ) -> Result { + validate_function_name(name)?; + validate_namespace(namespace)?; + self.object_job( + ExecuteQueryRequest { + query: format!( + "DESCRIBE FUNCTION {} WITH (version = {})", + quote_identifier(name), + literal(version) + ), + default_namespace_path: namespace.to_vec(), + parameters: None, + }, + definition, + ) + .await? + .wait() + .await + } + pub(in crate::remote) async fn list_functions( + &self, + namespace: &[String], + options: ListingOptions, + ) -> Result> { + validate_namespace(namespace)?; + let limit = options.page_limit.unwrap_or(100); + if limit == 0 || limit > i32::MAX as u32 { + return Err(Error::InvalidInput { + message: "page limit must be between 1 and i32::MAX".into(), + }); + } + let invalid = || Error::InvalidInput { + message: "invalid function page token for this database or namespace".into(), + }; + let after = if let Some(token) = options.page_token.filter(|token| !token.is_empty()) { + let bytes = URL_SAFE_NO_PAD.decode(token).map_err(|_| invalid())?; + let (database, prefix, path, name, version): ( + String, + Option, + Vec, + String, + String, + ) = serde_json::from_slice(&bytes).map_err(|_| invalid())?; + if database != self.inner.database + || prefix != self.inner.database_prefix + || path != namespace + { + return Err(invalid()); + } + Some((name, version)) + } else { + None + }; + let mut sql_options = vec![ + format!("limit = {}", u64::from(limit) + 1), + "include_definition = true".to_owned(), + ]; + if let Some((name, version)) = &after { + sql_options.push(format!("after_name = {}", literal(name))); + sql_options.push(format!("after_version = {}", literal(version))); + } + let mut items: Vec = self + .object_job( + ExecuteQueryRequest { + query: format!("SHOW FUNCTIONS WITH ({})", sql_options.join(", ")), + default_namespace_path: namespace.to_vec(), + parameters: None, + }, + |batches| { + strings(&batches, "definition")? + .into_iter() + .map(|value| { + serde_json::from_str(&value).map_err(|error| Error::Runtime { + message: error.to_string(), + }) + }) + .collect::>>() + }, + ) + .await? + .wait() + .await?; + items.sort_by(|a, b| (a.name(), a.version()).cmp(&(b.name(), b.version()))); + if let Some((name, version)) = after { + items.retain(|item| (item.name(), item.version()) > (name.as_str(), version.as_str())); + } + let page_token = if items.len() > limit as usize { + items.truncate(limit as usize); + let last = items.last().unwrap(); + let bytes = serde_json::to_vec(&( + &self.inner.database, + &self.inner.database_prefix, + namespace, + last.name(), + last.version(), + )) + .map_err(|error| Error::Runtime { + message: error.to_string(), + })?; + Some(URL_SAFE_NO_PAD.encode(bytes)) + } else { + None + }; + Ok(ListingPage { items, page_token }) + } + pub(in crate::remote) async fn drop_function( + &self, + name: &str, + version: &str, + namespace: &[String], + ) -> Result<(bool, Job)> { + validate_function_name(name)?; + validate_namespace(namespace)?; + let job = self + .object_job( + ExecuteQueryRequest { + query: format!( + "DROP FUNCTION IF EXISTS {} WITH (version = {})", + quote_identifier(name), + literal(version) + ), + default_namespace_path: namespace.to_vec(), + parameters: None, + }, + |batches| { + let rows = strings(&batches, "status")?; + match rows.as_slice() { + [status] if status == "dropped" => Ok(true), + [status] if status == "not_found" => Ok(false), + _ => Err(Error::Runtime { + message: "invalid DROP FUNCTION result".into(), + }), + } + }, + ) + .await?; + let id = job.id().unwrap().parse().map_err(|_| Error::Runtime { + message: "invalid SQL query job id".into(), + })?; + let dropped = job.wait().await?; + Ok(( + dropped, + Job::from_query_task(id, tokio::spawn(async { Ok(()) }), || async { Ok(()) }), + )) + } +} diff --git a/rust/lancedb/src/remote/sql/objects.rs b/rust/lancedb/src/remote/sql/objects.rs new file mode 100644 index 000000000..2c4d33614 --- /dev/null +++ b/rust/lancedb/src/remote/sql/objects.rs @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::SqlClient; +use crate::{ + database::ExecuteQueryRequest, + error::{Error, Result}, + job::Job, +}; +use arrow_array::{Array, RecordBatch, StringArray}; +use futures::TryStreamExt; +use serde::{Serialize, de::DeserializeOwned}; +use std::sync::Arc; + +pub(super) fn literal(value: &str) -> String { + format!("'{}'", value.replace('\'', "''")) +} +pub(super) fn strings(batches: &[RecordBatch], name: &str) -> Result> { + batches + .iter() + .map(|batch| { + let values = batch + .column_by_name(name) + .and_then(|column| column.as_any().downcast_ref::()) + .ok_or_else(|| Error::Runtime { + message: format!("SQL result is missing UTF-8 column '{name}'"), + })?; + if values.null_count() != 0 { + return Err(Error::Runtime { + message: format!("SQL result contains null '{name}'"), + }); + } + Ok(values + .iter() + .map(|value| value.unwrap().to_owned()) + .collect::>()) + }) + .collect::>>() + .map(|rows| rows.into_iter().flatten().collect()) +} +pub(super) fn definition(batches: Vec) -> Result { + let values = strings(&batches, "definition")?; + let [value]: [String; 1] = values.try_into().map_err(|_| Error::Runtime { + message: "SQL definition command must return exactly one row".into(), + })?; + serde_json::from_str(&value).map_err(|error| Error::Runtime { + message: format!("invalid SQL definition: {error}"), + }) +} +impl SqlClient { + pub(super) async fn object_job( + &self, + request: ExecuteQueryRequest, + decode: F, + ) -> Result> + where + T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static, + F: FnOnce(Vec) -> Result + Send + 'static, + { + let query = Arc::new(self.execute(request).await?); + let id = query.id(); + let execution = query.clone(); + let task = tokio::spawn(async move { + let reader = execution.reader().await?; + decode(reader.try_collect().await?) + }); + Ok(Job::from_query_task(id, task, move || { + let query = query.clone(); + async move { query.cancel().await } + })) + } +} diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 3418235bd..40b6c1b81 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -2270,3 +2270,167 @@ async fn storage_integration_methods_use_sql_with_namespace_and_typed_metadata() let _ = shutdown.send(()); } } + +#[tokio::test] +async fn function_creation_and_description_use_sql_query_handles() { + const VERSION: &str = include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json" + ); + const REQUEST: &str = include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json" + ); + for defer in [false, true] { + let mut request_json: serde_json::Value = serde_json::from_str(REQUEST).unwrap(); + request_json["secret_bindings"] = serde_json::json!([{ + "kind": "env", + "variable": "HF_TOKEN", + "secret_ref": {"name": "hf-prod"}, + }]); + let mut expected_request = request_json.clone(); + expected_request.as_object_mut().unwrap().remove("name"); + let request = + crate::function::FunctionRegistrationRequest::from_json(&request_json.to_string()) + .unwrap(); + let result = RecordBatch::try_from_iter(vec![( + "definition", + Arc::new(StringArray::from(vec![VERSION])) as ArrayRef, + )]) + .unwrap(); + let service = TestSqlService { + adaptive: true, + defer, + result, + ..Default::default() + }; + let statements = service.statements.clone(); + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|request| -> http::Response { + panic!("unexpected REST request: {}", request.url().path()); + }) + .with_sql_client(client.clone()); + let job = db + .create_function_async(request, &["team".into()]) + .await + .unwrap(); + assert_eq!(job.job_type().as_deref(), Some("query")); + let id = Uuid::parse_str(job.id().unwrap()).unwrap(); + assert_eq!(job.wait().await.unwrap().name(), "embed"); + assert_eq!(job.wait().await.unwrap().version(), "1"); + assert_eq!( + client.describe(id).await.unwrap().status, + QueryStatus::Finished + ); + assert_eq!( + db.get_function("embed", "1", &["team".into()]) + .await + .unwrap() + .name(), + "embed" + ); + let seen = statements.lock().unwrap(); + let encoded = seen[0] + .strip_prefix("CREATE FUNCTION \"normalize_score\" FROM '") + .unwrap() + .strip_suffix('\'') + .unwrap(); + let sent: serde_json::Value = serde_json::from_str(&encoded.replace("''", "'")).unwrap(); + assert_eq!(sent, expected_request); + assert_eq!(seen[1], "DESCRIBE FUNCTION \"embed\" WITH (version = '1')"); + shutdown.send(()).unwrap(); + } +} +#[tokio::test] +async fn function_drop_keeps_the_query_id_and_reports_catalog_removal() { + for (status, expected) in [("dropped", true), ("not_found", false)] { + let result = RecordBatch::try_from_iter(vec![( + "status", + Arc::new(StringArray::from(vec![status])) as ArrayRef, + )]) + .unwrap(); + let service = TestSqlService { + adaptive: true, + defer: true, + result, + ..Default::default() + }; + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|_| -> http::Response { + panic!("unexpected REST request"); + }) + .with_sql_client(client.clone()); + let (dropped, job) = db.drop_function_async("embed", "1", &[]).await.unwrap(); + assert_eq!(dropped, expected); + assert_eq!(job.job_type().as_deref(), Some("query")); + let id = Uuid::parse_str(job.id().unwrap()).unwrap(); + job.wait().await.unwrap(); + job.wait().await.unwrap(); + assert_eq!( + client.describe(id).await.unwrap().status, + QueryStatus::Finished + ); + shutdown.send(()).unwrap(); + } +} + +#[tokio::test] +async fn function_listing_uses_bounded_sql_and_scoped_version_cursors() { + const VERSION: &str = include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json" + ); + let mut second: serde_json::Value = serde_json::from_str(VERSION).unwrap(); + second["version"] = serde_json::json!("2"); + let second = second.to_string(); + let result = RecordBatch::try_from_iter(vec![( + "definition", + Arc::new(StringArray::from(vec![VERSION, second.as_str()])) as ArrayRef, + )]) + .unwrap(); + let service = TestSqlService { + adaptive: true, + defer: true, + result, + ..Default::default() + }; + let statements = service.statements.clone(); + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|_| -> http::Response { + panic!("unexpected REST request"); + }) + .with_sql_client(client); + let first = db + .list_functions( + &["team".into()], + crate::listing::ListingOptions::default().page_limit(1), + ) + .await + .unwrap(); + assert_eq!(first.items.len(), 1); + let token = first.page_token.unwrap(); + let second = db + .list_functions( + &["team".into()], + crate::listing::ListingOptions::default() + .page_limit(1) + .page_token(token.clone()), + ) + .await + .unwrap(); + assert_eq!(second.items.len(), 1); + assert_eq!(second.items[0].version(), "2"); + assert!(second.page_token.is_none()); + assert!( + db.list_functions( + &["other".into()], + crate::listing::ListingOptions::default() + .page_limit(1) + .page_token(token) + ) + .await + .is_err() + ); + let seen = statements.lock().unwrap(); + assert_eq!(seen.len(), 2); + assert!(seen[0].contains("limit = 2")); + assert!(seen[1].contains("after_version = '1'")); + shutdown.send(()).unwrap(); +}