statement ok
CREATE SEQUENCE seq;

statement ok
CREATE TYPE weekday AS ENUM ('monday', 'tuesday', 'wednesday', 'thursday', 'friday');

statement ok
CREATE TABLE t_rewrite (
  v INT DEFAULT 0,
  w weekday DEFAULT 'monday'::weekday
);

# Need to turn declarative schema changer off because function `get_body_str`
# created below would resolve a descriptorless public schema "system.public"
# which is not supported in declarative schema changer. Declarative schema
# changer falls back to legacy schema changer, and the descriptor id counter is
# increased twice. It cause the test to fail due to id inconsistency.
skipif config local-legacy-schema-changer
statement ok
SET use_declarative_schema_changer = 'off'

statement ok
CREATE FUNCTION get_body_str(fn_name STRING) RETURNS STRING
LANGUAGE SQL
AS $$
  SELECT crdb_internal.pb_to_json(
    'cockroach.sql.sqlbase.Descriptor', descriptor, false
  )->'function'->'functionBody'
  FROM system.descriptor WHERE id = fn_name::regproc::int - 100000;
$$;

skipif config local-legacy-schema-changer
statement ok
SET use_declarative_schema_changer = 'on'

subtest rewrite_sql

statement ok
CREATE FUNCTION f_rewrite() RETURNS INT AS
$$
  SELECT nextval('seq');
$$ LANGUAGE SQL

query T
SELECT get_body_str('f_rewrite');
----
"SELECT nextval(106:::REGCLASS);"

statement ok
DROP FUNCTION f_rewrite();

statement ok
CREATE FUNCTION f_rewrite() RETURNS INT AS
$$
  INSERT INTO t_rewrite(v) VALUES (nextval('seq')) RETURNING v;
$$ LANGUAGE SQL

query T
SELECT get_body_str('f_rewrite');
----
"INSERT INTO test.public.t_rewrite(v) VALUES (nextval(106:::REGCLASS)) RETURNING v;"

statement ok
DROP FUNCTION f_rewrite();

statement ok
CREATE FUNCTION f_rewrite() RETURNS weekday AS
$$
  SELECT 'wednesday'::weekday;
$$ LANGUAGE SQL

query T
SELECT get_body_str('f_rewrite');
----
"SELECT b'\\x80':::@100107;"

statement ok
DROP FUNCTION f_rewrite();

statement ok
CREATE FUNCTION f_rewrite() RETURNS weekday AS
$$
  UPDATE t_rewrite SET w = 'thursday'::weekday WHERE w = 'wednesday'::weekday RETURNING w;
$$ LANGUAGE SQL

query T
SELECT get_body_str('f_rewrite');
----
"UPDATE test.public.t_rewrite SET w = b'\\xa0':::@100107 WHERE w = b'\\x80':::@100107 RETURNING w;"

statement ok
DROP FUNCTION f_rewrite();

statement ok
CREATE FUNCTION f_rewrite(OUT weekday) AS
$$
  SELECT 'thursday'::weekday;
$$ LANGUAGE SQL

query T
SELECT get_body_str('f_rewrite');
----
"SELECT b'\\xa0':::@100107;"

statement ok
DROP FUNCTION f_rewrite();

statement ok
CREATE FUNCTION f_rewrite(INOUT a weekday) AS
$$
  SELECT 'thursday'::weekday;
$$ LANGUAGE SQL

query T
SELECT get_body_str('f_rewrite');
----
"SELECT b'\\xa0':::@100107;"

statement ok
DROP FUNCTION f_rewrite;

subtest end

subtest rewrite_proc

statement ok
CREATE PROCEDURE p_rewrite() AS
$$
  INSERT INTO t_rewrite(v) VALUES (nextval('seq')) RETURNING v;
$$ LANGUAGE SQL

query T
SELECT get_body_str('p_rewrite');
----
"INSERT INTO test.public.t_rewrite(v) VALUES (nextval(106:::REGCLASS)) RETURNING v;"

statement ok
DROP PROCEDURE p_rewrite();

statement ok
CREATE PROCEDURE p_rewrite() AS
$$
  UPDATE t_rewrite SET w = 'thursday'::weekday WHERE w = 'wednesday'::weekday RETURNING w;
$$ LANGUAGE SQL

query T
SELECT get_body_str('p_rewrite');
----
"UPDATE test.public.t_rewrite SET w = b'\\xa0':::@100107 WHERE w = b'\\x80':::@100107 RETURNING w;"

statement ok
DROP PROCEDURE p_rewrite();

statement ok
CREATE PROCEDURE p_rewrite(OUT weekday) AS
$$
  SELECT 'thursday'::weekday;
$$ LANGUAGE SQL

query T
SELECT get_body_str('p_rewrite');
----
"SELECT b'\\xa0':::@100107;"

statement ok
DROP PROCEDURE p_rewrite();

statement ok
CREATE PROCEDURE p_rewrite(INOUT a weekday) AS
$$
  SELECT 'thursday'::weekday;
$$ LANGUAGE SQL

query T
SELECT get_body_str('p_rewrite');
----
"SELECT b'\\xa0':::@100107;"

statement ok
DROP PROCEDURE p_rewrite;

subtest end

subtest rewrite_udf_calling_duf

statement ok
CREATE FUNCTION nested_func() RETURNS INT AS $$
  SELECT 1;
$$ LANGUAGE SQL

statement ok
CREATE PROCEDURE p_rewrite() AS $$
  SELECT nested_func();
  SELECT * FROM nested_func();
$$ LANGUAGE SQL

# TODO(fqazi): Renaming function calls will break today until #120351 is completed.
query T
SELECT get_body_str('p_rewrite');
----
"SELECT public.nested_func();\nSELECT nested_func FROM ROWS FROM (public.nested_func());"

subtest end
