You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "kaxil (via GitHub)" <gi...@apache.org> on 2024/10/22 01:26:45 UTC
[PR] Make Task Instance primary key be a UUID [airflow]
kaxil opened a new pull request, #43243:
URL: https://github.com/apache/airflow/pull/43243
[skip ci]
closes https://github.com/apache/airflow/issues/43161 part of [AIP-72](https://github.com/orgs/apache/projects/405) | [AIP-72 details here](https://cwiki.apache.org/confluence/display/AIRFLOW/AIP-72+Task+Execution+Interface+aka+Task+SDK)
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [newsfragments](https://github.com/apache/airflow/tree/main/newsfragments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818885052
##########
airflow/models/taskinstance.py:
##########
@@ -1813,10 +1817,16 @@ class TaskInstance(Base, LoggingMixin):
"""
__tablename__ = "task_instance"
- task_id = Column(StringID(), primary_key=True, nullable=False)
- dag_id = Column(StringID(), primary_key=True, nullable=False)
- run_id = Column(StringID(), primary_key=True, nullable=False)
- map_index = Column(Integer, primary_key=True, nullable=False, server_default=text("-1"))
+ id = Column(
+ String(36).with_variant(postgresql.UUID(as_uuid=False), "postgresql"),
+ primary_key=True,
+ default=lambda: str(uuid6.uuid7()),
Review Comment:
Maybe we don't need to force it to a string but instead add the SQLA typeadapter https://github.com/apache/airflow/blob/8a6b1fc4e0ab8ba42cc47d8390fcd62a41e924dd/airflow/utils/sqlalchemy.py#L48
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #43243:
URL: https://github.com/apache/airflow/pull/43243#issuecomment-2441350388
I do think long-term/by release of 3.0 TaskInstanceHistory should also have a uuid.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818881363
##########
airflow/migrations/versions/0041_3_0_0_add_uuid_primary_key_to_task_instance_.py:
##########
@@ -0,0 +1,273 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add UUID primary key to ``task_instance`` table.
+
+Revision ID: d59cbbef95eb
+Revises: 05234396c6fc
+Create Date: 2024-10-21 22:39:12.394079
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import text
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "d59cbbef95eb"
+down_revision = "05234396c6fc"
+branch_labels = "None"
+depends_on = None
+airflow_version = "3.0.0"
+
+
+# PostgreSQL-specific UUID v7 function
+pg_uuid7_fn = """
+DO $$
+DECLARE
+ pgcrypto_installed BOOLEAN;
+BEGIN
+ -- Check if pgcrypto is already installed
+ pgcrypto_installed := EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pgcrypto');
+
+ -- Attempt to create pgcrypto if it is not installed
+ IF NOT pgcrypto_installed THEN
+ BEGIN
+ CREATE EXTENSION pgcrypto;
+ pgcrypto_installed := TRUE;
+ RAISE NOTICE 'pgcrypto extension successfully created.';
+ EXCEPTION
+ WHEN insufficient_privilege THEN
+ RAISE NOTICE 'pgcrypto extension could not be installed due to insufficient privileges; using fallback';
+ pgcrypto_installed := FALSE;
+ WHEN OTHERS THEN
+ RAISE NOTICE 'An unexpected error occurred while attempting to install pgcrypto; using fallback';
+ pgcrypto_installed := FALSE;
+ END;
+ END IF;
+END $$;
+
+CREATE OR REPLACE FUNCTION uuid_generate_v7(p_timestamp timestamp with time zone)
+RETURNS uuid
+LANGUAGE plpgsql
+PARALLEL SAFE
+AS $$
+DECLARE
+ unix_time_ms CONSTANT bytea NOT NULL DEFAULT substring(int8send((extract(epoch FROM p_timestamp) * 1000)::bigint) from 3);
+ buffer bytea;
+ pgcrypto_installed BOOLEAN := EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pgcrypto');
+BEGIN
+ -- Use pgcrypto if available, otherwise use the fallback
+ -- fallback from https://brandur.org/fragments/secure-bytes-without-pgcrypto
+ IF pgcrypto_installed THEN
+ buffer := unix_time_ms || gen_random_bytes(10);
+ ELSE
+ buffer := unix_time_ms || substring(uuid_send(gen_random_uuid()) FROM 1 FOR 5) ||
+ substring(uuid_send(gen_random_uuid()) FROM 12 FOR 5);
+ END IF;
+
+ -- Set UUID version and variant bits
+ buffer := set_byte(buffer, 6, (b'0111' || get_byte(buffer, 6)::bit(4))::bit(8)::int);
+ buffer := set_byte(buffer, 8, (b'10' || get_byte(buffer, 8)::bit(6))::bit(8)::int);
+ RETURN encode(buffer, 'hex')::uuid;
+END
+$$;
+"""
+
+pg_uuid7_fn_drop = """
+DROP FUNCTION IF EXISTS uuid_generate_v7(timestamp with time zone);
+"""
+
+# MySQL-specific UUID v7 function
+mysql_uuid7_fn = """
+DROP FUNCTION IF EXISTS uuid_generate_v7;
+CREATE FUNCTION uuid_generate_v7(p_timestamp DATETIME(3))
+RETURNS CHAR(36)
+DETERMINISTIC
+BEGIN
+ DECLARE unix_time_ms BIGINT;
+ DECLARE time_hex CHAR(12);
+ DECLARE rand_hex CHAR(24);
+ DECLARE uuid CHAR(36);
+
+ -- Convert the passed timestamp to milliseconds since epoch
+ SET unix_time_ms = UNIX_TIMESTAMP(p_timestamp) * 1000;
+ SET time_hex = LPAD(HEX(unix_time_ms), 12, '0');
+ SET rand_hex = CONCAT(
+ LPAD(HEX(FLOOR(RAND() * POW(2,32))), 8, '0'),
+ LPAD(HEX(FLOOR(RAND() * POW(2,32))), 8, '0')
+ );
+ SET rand_hex = CONCAT(SUBSTRING(rand_hex, 1, 4), '7', SUBSTRING(rand_hex, 6));
+ SET rand_hex = CONCAT(SUBSTRING(rand_hex, 1, 12), '8', SUBSTRING(rand_hex, 14));
+
+ SET uuid = LOWER(CONCAT(
+ SUBSTRING(time_hex, 1, 8), '-',
+ SUBSTRING(time_hex, 9, 4), '-',
+ SUBSTRING(rand_hex, 1, 4), '-',
+ SUBSTRING(rand_hex, 5, 4), '-',
+ SUBSTRING(rand_hex, 9)
+ ));
+
+ RETURN uuid;
+END;
+"""
+
+mysql_uuid7_fn_drop = """
+DROP FUNCTION IF EXISTS uuid_generate_v7;
+"""
+
+ti_table = "task_instance"
+
+# Foreign key columns from task_instance
+ti_fk_cols = ["dag_id", "task_id", "run_id", "map_index"]
+
+# Foreign key constraints from other tables to task_instance
+ti_fk_constrains = [
Review Comment:
Done in https://github.com/apache/airflow/pull/43243/commits/0f5dc0a0e496a38d53003ee9dc2941af0d8c9b74
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make Task Instance primary key be a UUID [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #43243:
URL: https://github.com/apache/airflow/pull/43243#issuecomment-2439774584
After updating the primary key on the `TaskInstance` model to a UUID7 `id` field (from the composite primary key of `["dag_id", "task_id", "run_id", "map_index"]`), I have now run into issues with using `session.merge()` in our code (especially tests!).
The `session.merge()` function operates strictly on primary keys as mentioned in [this SQLAlchemy docs](https://docs.sqlalchemy.org/en/20/orm/session_state_management.html#merging), which means it no longer recognizes the unique constraint on `["dag_id", "task_id", "run_id", "map_index"]` to identify existing `TaskInstance` records. This leads to issues in cases where `session.merge()` was previously used to either update or insert TaskInstance records, as it now fails to locate the intended record by the unique constraint.
Example:
https://github.com/apache/airflow/blob/b2ecb6c9caacd2ca7a2c1519ca892b19b2bb35fb/tests/models/test_dag.py#L2013-L2026
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #43243:
URL: https://github.com/apache/airflow/pull/43243#issuecomment-2441409245
One problem with generating the unique hash with other fields, is that it will be even **more** difficult to add new fields, so probably getting rid of session.merge() everywhere and replacing it with our own way of merging objects and database entries by retrieving the unique id based on the unique fields from the DB before merging is likely a long-term better approach.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #43243:
URL: https://github.com/apache/airflow/pull/43243#issuecomment-2441490772
> UUID v7 is explicitly designed to support distributed databases with high insert rates due to its temporal ordering.
Yep. that settles it as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818240774
##########
tests/jobs/test_local_task_job.py:
##########
@@ -951,11 +953,21 @@ def test_mini_scheduler_works_with_wait_for_upstream(self, caplog, get_test_dag)
task_k = dag.get_task("K")
task_l = dag.get_task("L")
with create_session() as session:
- ti_k = TaskInstance(task_k, run_id=dr.run_id, state=State.SUCCESS)
- ti_b = TaskInstance(task_l, run_id=dr.run_id, state=State.SUCCESS)
+ ti_k = dr.get_task_instance(task_k.task_id, session=session)
Review Comment:
The reason for this and other changes in this file is explained in [**Issues identified Point (1)**](https://github.com/apache/airflow/pull/43243#issuecomment-2439774584)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818880931
##########
airflow/models/taskinstance.py:
##########
@@ -1813,10 +1817,16 @@ class TaskInstance(Base, LoggingMixin):
"""
__tablename__ = "task_instance"
- task_id = Column(StringID(), primary_key=True, nullable=False)
- dag_id = Column(StringID(), primary_key=True, nullable=False)
- run_id = Column(StringID(), primary_key=True, nullable=False)
- map_index = Column(Integer, primary_key=True, nullable=False, server_default=text("-1"))
+ id = Column(
+ String(36).with_variant(postgresql.UUID(as_uuid=False), "postgresql"),
+ primary_key=True,
+ default=lambda: str(uuid6.uuid7()),
Review Comment:
Good point. Done in https://github.com/apache/airflow/pull/43243/commits/0f5dc0a0e496a38d53003ee9dc2941af0d8c9b74
##########
airflow/migrations/versions/0041_3_0_0_add_uuid_primary_key_to_task_instance_.py:
##########
@@ -0,0 +1,273 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add UUID primary key to ``task_instance`` table.
+
+Revision ID: d59cbbef95eb
+Revises: 05234396c6fc
+Create Date: 2024-10-21 22:39:12.394079
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import text
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "d59cbbef95eb"
+down_revision = "05234396c6fc"
+branch_labels = "None"
+depends_on = None
+airflow_version = "3.0.0"
+
+
+# PostgreSQL-specific UUID v7 function
+pg_uuid7_fn = """
+DO $$
+DECLARE
+ pgcrypto_installed BOOLEAN;
+BEGIN
+ -- Check if pgcrypto is already installed
+ pgcrypto_installed := EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pgcrypto');
+
+ -- Attempt to create pgcrypto if it is not installed
+ IF NOT pgcrypto_installed THEN
+ BEGIN
+ CREATE EXTENSION pgcrypto;
+ pgcrypto_installed := TRUE;
+ RAISE NOTICE 'pgcrypto extension successfully created.';
+ EXCEPTION
+ WHEN insufficient_privilege THEN
+ RAISE NOTICE 'pgcrypto extension could not be installed due to insufficient privileges; using fallback';
+ pgcrypto_installed := FALSE;
+ WHEN OTHERS THEN
+ RAISE NOTICE 'An unexpected error occurred while attempting to install pgcrypto; using fallback';
+ pgcrypto_installed := FALSE;
+ END;
+ END IF;
+END $$;
+
+CREATE OR REPLACE FUNCTION uuid_generate_v7(p_timestamp timestamp with time zone)
+RETURNS uuid
+LANGUAGE plpgsql
+PARALLEL SAFE
+AS $$
+DECLARE
+ unix_time_ms CONSTANT bytea NOT NULL DEFAULT substring(int8send((extract(epoch FROM p_timestamp) * 1000)::bigint) from 3);
+ buffer bytea;
+ pgcrypto_installed BOOLEAN := EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pgcrypto');
+BEGIN
+ -- Use pgcrypto if available, otherwise use the fallback
+ -- fallback from https://brandur.org/fragments/secure-bytes-without-pgcrypto
+ IF pgcrypto_installed THEN
+ buffer := unix_time_ms || gen_random_bytes(10);
+ ELSE
+ buffer := unix_time_ms || substring(uuid_send(gen_random_uuid()) FROM 1 FOR 5) ||
+ substring(uuid_send(gen_random_uuid()) FROM 12 FOR 5);
+ END IF;
+
+ -- Set UUID version and variant bits
+ buffer := set_byte(buffer, 6, (b'0111' || get_byte(buffer, 6)::bit(4))::bit(8)::int);
+ buffer := set_byte(buffer, 8, (b'10' || get_byte(buffer, 8)::bit(6))::bit(8)::int);
+ RETURN encode(buffer, 'hex')::uuid;
+END
+$$;
+"""
+
+pg_uuid7_fn_drop = """
+DROP FUNCTION IF EXISTS uuid_generate_v7(timestamp with time zone);
+"""
+
+# MySQL-specific UUID v7 function
+mysql_uuid7_fn = """
+DROP FUNCTION IF EXISTS uuid_generate_v7;
+CREATE FUNCTION uuid_generate_v7(p_timestamp DATETIME(3))
+RETURNS CHAR(36)
+DETERMINISTIC
+BEGIN
+ DECLARE unix_time_ms BIGINT;
+ DECLARE time_hex CHAR(12);
+ DECLARE rand_hex CHAR(24);
+ DECLARE uuid CHAR(36);
+
+ -- Convert the passed timestamp to milliseconds since epoch
+ SET unix_time_ms = UNIX_TIMESTAMP(p_timestamp) * 1000;
+ SET time_hex = LPAD(HEX(unix_time_ms), 12, '0');
+ SET rand_hex = CONCAT(
+ LPAD(HEX(FLOOR(RAND() * POW(2,32))), 8, '0'),
+ LPAD(HEX(FLOOR(RAND() * POW(2,32))), 8, '0')
+ );
+ SET rand_hex = CONCAT(SUBSTRING(rand_hex, 1, 4), '7', SUBSTRING(rand_hex, 6));
+ SET rand_hex = CONCAT(SUBSTRING(rand_hex, 1, 12), '8', SUBSTRING(rand_hex, 14));
+
+ SET uuid = LOWER(CONCAT(
+ SUBSTRING(time_hex, 1, 8), '-',
+ SUBSTRING(time_hex, 9, 4), '-',
+ SUBSTRING(rand_hex, 1, 4), '-',
+ SUBSTRING(rand_hex, 5, 4), '-',
+ SUBSTRING(rand_hex, 9)
+ ));
+
+ RETURN uuid;
+END;
+"""
+
+mysql_uuid7_fn_drop = """
+DROP FUNCTION IF EXISTS uuid_generate_v7;
+"""
+
+ti_table = "task_instance"
+
+# Foreign key columns from task_instance
+ti_fk_cols = ["dag_id", "task_id", "run_id", "map_index"]
+
+# Foreign key constraints from other tables to task_instance
+ti_fk_constrains = [
+ {"table": "rendered_task_instance_fields", "fk": "rtif_ti_fkey"},
+ {"table": "task_fail", "fk": "task_fail_ti_fkey"},
+ {"table": "task_instance_history", "fk": "task_instance_history_ti_fkey"},
+ {"table": "task_instance_note", "fk": "task_instance_note_ti_fkey"},
+ {"table": "task_map", "fk": "task_map_task_instance_fkey"},
+ {"table": "task_reschedule", "fk": "task_reschedule_ti_fkey"},
+ {"table": "xcom", "fk": "xcom_task_instance_fkey"},
+]
+
+
+def _get_type_id_column(dialect_name: str) -> sa.types.TypeEngine:
+ # For PostgreSQL, use the UUID type directly as it is more efficient
+ if dialect_name == "postgresql":
+ return postgresql.UUID(as_uuid=False)
+ # For other databases, use String(36) to match UUID format
+ else:
+ return sa.String(36)
+
+
+def upgrade():
+ """Add UUID primary key to task instance table."""
+ conn = op.get_bind()
+ dialect_name = conn.dialect.name
+
+ op.add_column("task_instance", sa.Column("id", _get_type_id_column(dialect_name), nullable=True))
+
+ if dialect_name == "postgresql":
+ op.execute(pg_uuid7_fn)
+
+ # TODO: Add batching to handle updates in smaller chunks for large tables to avoid locking
+ # Migrate existing rows with UUID v7 using a timestamp-based generation
+ op.execute(
+ "UPDATE task_instance SET id = uuid_generate_v7(coalesce(queued_dttm, start_date, clock_timestamp()))"
+ )
+
+ op.execute(pg_uuid7_fn_drop)
+
+ # Drop existing primary key constraint to task_instance table
+ op.execute("ALTER TABLE IF EXISTS task_instance DROP CONSTRAINT task_instance_pkey CASCADE")
+
+ elif dialect_name == "mysql":
+ op.execute(mysql_uuid7_fn)
+
+ # Migrate existing rows with UUID v7
+ op.execute("""
+ UPDATE task_instance
+ SET id = uuid_generate_v7(coalesce(queued_dttm, start_date, NOW(3)))
+ WHERE id IS NULL
+ """)
+
+ op.execute(mysql_uuid7_fn_drop)
Review Comment:
Good point. Done in https://github.com/apache/airflow/pull/43243/commits/0f5dc0a0e496a38d53003ee9dc2941af0d8c9b74
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818239964
##########
airflow/api_connexion/endpoints/task_instance_endpoint.py:
##########
@@ -533,9 +533,11 @@ def post_set_task_instances_state(*, dag_id: str, session: Session = NEW_SESSION
detail=f"Task instance not found for task {task_id!r} on execution_date {execution_date}"
)
- if run_id and not session.get(
- TI, {"task_id": task_id, "dag_id": dag_id, "run_id": run_id, "map_index": -1}
- ):
+ select_stmt = select(TI).where(
+ TI.dag_id == dag_id, TI.task_id == task_id, TI.run_id == run_id, TI.map_index == -1
+ )
+
+ if run_id and not session.scalars(select_stmt).one_or_none():
Review Comment:
The reason for this is explained in [**Issues identified Point (3)**](https://github.com/apache/airflow/pull/43243#issuecomment-2439774584)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #43243:
URL: https://github.com/apache/airflow/pull/43243#issuecomment-2441393201
> Yeah, our use of `session.merge()` in Airflow is v. bad :(
Agree. This is one of the reasons I wanted to avoid all those db-schema changes for the initial proposal of multi-team and proposed the team-id prefix, becaue I was afraid we are going to have similar ripple-effects that will result in rewriting big part of our database code. Which seem that it was a pretty justified worry.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818240481
##########
providers/tests/openlineage/extractors/test_manager.py:
##########
@@ -313,6 +313,7 @@ def use_read():
dr = dag_maker.create_dagrun()
ti = TaskInstance(task=task, run_id=dr.run_id)
+ ti.refresh_from_db()
Review Comment:
The reason for this is explained in [**Issues identified Point (1)**](https://github.com/apache/airflow/pull/43243#issuecomment-2439774584)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #43243:
URL: https://github.com/apache/airflow/pull/43243#issuecomment-2441487051
@potiuk Yeah, I considered hashing but it is bad for Databases for indexing and, as a result, querying since it won't have temporal properties apart from Collision Handling complexity.
UUID v7 is explicitly designed to support distributed databases with high insert rates due to its temporal ordering.
> BTW. One of the ways it could be helped - we **could** potentially generate the unique id by hashing the remaining fields - with appropriate hashing algorithm, we could have very low probability of collision (and maybe we could implement a mechanism to detect collisions and implement handling of those collision in similar ways they are handled in hash maps) - that would be another way how we could approach it (and there the merge code could be simpler and still use session.merge()
>
> But this one also comes with it's own set of difficulties and collision handling is likely going to be complex. But I thought it's worth menioning it here as an option.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #43243:
URL: https://github.com/apache/airflow/pull/43243#issuecomment-2441487328
> Yeah, our use of `session.merge()` in Airflow is v. bad :(
Indeed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818989737
##########
tests/jobs/test_local_task_job.py:
##########
@@ -902,7 +902,9 @@ def test_fast_follow(
ti_by_task_id = {}
with create_session() as session:
for task_id in init_state:
- ti = TaskInstance(dag.get_task(task_id), run_id=dag_run.run_id, state=init_state[task_id])
+ ti = TaskInstance(dag.get_task(task_id), run_id=dag_run.run_id)
+ ti.refresh_from_db()
Review Comment:
That feels a weird one to be added to the constructor though. At some point, we should just do a better Session & State management in Airflow with the `provide_session`, `session.merge` shenanigans.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make Task Instance primary key be a UUID [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil closed pull request #43243: Make Task Instance primary key be a UUID
URL: https://github.com/apache/airflow/pull/43243
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "ianbuss (via GitHub)" <gi...@apache.org>.
ianbuss commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818652707
##########
airflow/migrations/versions/0041_3_0_0_add_uuid_primary_key_to_task_instance_.py:
##########
@@ -0,0 +1,273 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add UUID primary key to ``task_instance`` table.
+
+Revision ID: d59cbbef95eb
+Revises: 05234396c6fc
+Create Date: 2024-10-21 22:39:12.394079
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import text
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "d59cbbef95eb"
+down_revision = "05234396c6fc"
+branch_labels = "None"
+depends_on = None
+airflow_version = "3.0.0"
+
+
+# PostgreSQL-specific UUID v7 function
+pg_uuid7_fn = """
+DO $$
+DECLARE
+ pgcrypto_installed BOOLEAN;
+BEGIN
+ -- Check if pgcrypto is already installed
+ pgcrypto_installed := EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pgcrypto');
+
+ -- Attempt to create pgcrypto if it is not installed
+ IF NOT pgcrypto_installed THEN
+ BEGIN
+ CREATE EXTENSION pgcrypto;
+ pgcrypto_installed := TRUE;
+ RAISE NOTICE 'pgcrypto extension successfully created.';
+ EXCEPTION
+ WHEN insufficient_privilege THEN
+ RAISE NOTICE 'pgcrypto extension could not be installed due to insufficient privileges; using fallback';
+ pgcrypto_installed := FALSE;
+ WHEN OTHERS THEN
+ RAISE NOTICE 'An unexpected error occurred while attempting to install pgcrypto; using fallback';
+ pgcrypto_installed := FALSE;
+ END;
+ END IF;
+END $$;
+
+CREATE OR REPLACE FUNCTION uuid_generate_v7(p_timestamp timestamp with time zone)
+RETURNS uuid
+LANGUAGE plpgsql
+PARALLEL SAFE
+AS $$
+DECLARE
+ unix_time_ms CONSTANT bytea NOT NULL DEFAULT substring(int8send((extract(epoch FROM p_timestamp) * 1000)::bigint) from 3);
+ buffer bytea;
+ pgcrypto_installed BOOLEAN := EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pgcrypto');
+BEGIN
+ -- Use pgcrypto if available, otherwise use the fallback
+ -- fallback from https://brandur.org/fragments/secure-bytes-without-pgcrypto
+ IF pgcrypto_installed THEN
+ buffer := unix_time_ms || gen_random_bytes(10);
+ ELSE
+ buffer := unix_time_ms || substring(uuid_send(gen_random_uuid()) FROM 1 FOR 5) ||
+ substring(uuid_send(gen_random_uuid()) FROM 12 FOR 5);
+ END IF;
+
+ -- Set UUID version and variant bits
+ buffer := set_byte(buffer, 6, (b'0111' || get_byte(buffer, 6)::bit(4))::bit(8)::int);
+ buffer := set_byte(buffer, 8, (b'10' || get_byte(buffer, 8)::bit(6))::bit(8)::int);
+ RETURN encode(buffer, 'hex')::uuid;
+END
+$$;
+"""
+
+pg_uuid7_fn_drop = """
+DROP FUNCTION IF EXISTS uuid_generate_v7(timestamp with time zone);
+"""
+
+# MySQL-specific UUID v7 function
+mysql_uuid7_fn = """
+DROP FUNCTION IF EXISTS uuid_generate_v7;
+CREATE FUNCTION uuid_generate_v7(p_timestamp DATETIME(3))
+RETURNS CHAR(36)
+DETERMINISTIC
+BEGIN
+ DECLARE unix_time_ms BIGINT;
+ DECLARE time_hex CHAR(12);
+ DECLARE rand_hex CHAR(24);
+ DECLARE uuid CHAR(36);
+
+ -- Convert the passed timestamp to milliseconds since epoch
+ SET unix_time_ms = UNIX_TIMESTAMP(p_timestamp) * 1000;
+ SET time_hex = LPAD(HEX(unix_time_ms), 12, '0');
+ SET rand_hex = CONCAT(
+ LPAD(HEX(FLOOR(RAND() * POW(2,32))), 8, '0'),
+ LPAD(HEX(FLOOR(RAND() * POW(2,32))), 8, '0')
+ );
+ SET rand_hex = CONCAT(SUBSTRING(rand_hex, 1, 4), '7', SUBSTRING(rand_hex, 6));
+ SET rand_hex = CONCAT(SUBSTRING(rand_hex, 1, 12), '8', SUBSTRING(rand_hex, 14));
+
+ SET uuid = LOWER(CONCAT(
+ SUBSTRING(time_hex, 1, 8), '-',
+ SUBSTRING(time_hex, 9, 4), '-',
+ SUBSTRING(rand_hex, 1, 4), '-',
+ SUBSTRING(rand_hex, 5, 4), '-',
+ SUBSTRING(rand_hex, 9)
+ ));
+
+ RETURN uuid;
+END;
+"""
+
+mysql_uuid7_fn_drop = """
+DROP FUNCTION IF EXISTS uuid_generate_v7;
+"""
+
+ti_table = "task_instance"
+
+# Foreign key columns from task_instance
+ti_fk_cols = ["dag_id", "task_id", "run_id", "map_index"]
+
+# Foreign key constraints from other tables to task_instance
+ti_fk_constrains = [
Review Comment:
```suggestion
ti_fk_constraints = [
```
nit
##########
airflow/migrations/versions/0041_3_0_0_add_uuid_primary_key_to_task_instance_.py:
##########
@@ -0,0 +1,273 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add UUID primary key to ``task_instance`` table.
+
+Revision ID: d59cbbef95eb
+Revises: 05234396c6fc
+Create Date: 2024-10-21 22:39:12.394079
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import text
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "d59cbbef95eb"
+down_revision = "05234396c6fc"
+branch_labels = "None"
+depends_on = None
+airflow_version = "3.0.0"
+
+
+# PostgreSQL-specific UUID v7 function
+pg_uuid7_fn = """
+DO $$
+DECLARE
+ pgcrypto_installed BOOLEAN;
+BEGIN
+ -- Check if pgcrypto is already installed
+ pgcrypto_installed := EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pgcrypto');
+
+ -- Attempt to create pgcrypto if it is not installed
+ IF NOT pgcrypto_installed THEN
+ BEGIN
+ CREATE EXTENSION pgcrypto;
+ pgcrypto_installed := TRUE;
+ RAISE NOTICE 'pgcrypto extension successfully created.';
+ EXCEPTION
+ WHEN insufficient_privilege THEN
+ RAISE NOTICE 'pgcrypto extension could not be installed due to insufficient privileges; using fallback';
+ pgcrypto_installed := FALSE;
+ WHEN OTHERS THEN
+ RAISE NOTICE 'An unexpected error occurred while attempting to install pgcrypto; using fallback';
+ pgcrypto_installed := FALSE;
+ END;
+ END IF;
+END $$;
+
+CREATE OR REPLACE FUNCTION uuid_generate_v7(p_timestamp timestamp with time zone)
+RETURNS uuid
+LANGUAGE plpgsql
+PARALLEL SAFE
+AS $$
+DECLARE
+ unix_time_ms CONSTANT bytea NOT NULL DEFAULT substring(int8send((extract(epoch FROM p_timestamp) * 1000)::bigint) from 3);
+ buffer bytea;
+ pgcrypto_installed BOOLEAN := EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pgcrypto');
+BEGIN
+ -- Use pgcrypto if available, otherwise use the fallback
+ -- fallback from https://brandur.org/fragments/secure-bytes-without-pgcrypto
+ IF pgcrypto_installed THEN
+ buffer := unix_time_ms || gen_random_bytes(10);
+ ELSE
+ buffer := unix_time_ms || substring(uuid_send(gen_random_uuid()) FROM 1 FOR 5) ||
+ substring(uuid_send(gen_random_uuid()) FROM 12 FOR 5);
+ END IF;
+
+ -- Set UUID version and variant bits
+ buffer := set_byte(buffer, 6, (b'0111' || get_byte(buffer, 6)::bit(4))::bit(8)::int);
+ buffer := set_byte(buffer, 8, (b'10' || get_byte(buffer, 8)::bit(6))::bit(8)::int);
+ RETURN encode(buffer, 'hex')::uuid;
+END
+$$;
+"""
+
+pg_uuid7_fn_drop = """
+DROP FUNCTION IF EXISTS uuid_generate_v7(timestamp with time zone);
+"""
+
+# MySQL-specific UUID v7 function
+mysql_uuid7_fn = """
+DROP FUNCTION IF EXISTS uuid_generate_v7;
+CREATE FUNCTION uuid_generate_v7(p_timestamp DATETIME(3))
+RETURNS CHAR(36)
+DETERMINISTIC
+BEGIN
+ DECLARE unix_time_ms BIGINT;
+ DECLARE time_hex CHAR(12);
+ DECLARE rand_hex CHAR(24);
+ DECLARE uuid CHAR(36);
+
+ -- Convert the passed timestamp to milliseconds since epoch
+ SET unix_time_ms = UNIX_TIMESTAMP(p_timestamp) * 1000;
+ SET time_hex = LPAD(HEX(unix_time_ms), 12, '0');
+ SET rand_hex = CONCAT(
+ LPAD(HEX(FLOOR(RAND() * POW(2,32))), 8, '0'),
+ LPAD(HEX(FLOOR(RAND() * POW(2,32))), 8, '0')
+ );
+ SET rand_hex = CONCAT(SUBSTRING(rand_hex, 1, 4), '7', SUBSTRING(rand_hex, 6));
+ SET rand_hex = CONCAT(SUBSTRING(rand_hex, 1, 12), '8', SUBSTRING(rand_hex, 14));
+
+ SET uuid = LOWER(CONCAT(
+ SUBSTRING(time_hex, 1, 8), '-',
+ SUBSTRING(time_hex, 9, 4), '-',
+ SUBSTRING(rand_hex, 1, 4), '-',
+ SUBSTRING(rand_hex, 5, 4), '-',
+ SUBSTRING(rand_hex, 9)
+ ));
+
+ RETURN uuid;
+END;
+"""
+
+mysql_uuid7_fn_drop = """
+DROP FUNCTION IF EXISTS uuid_generate_v7;
+"""
+
+ti_table = "task_instance"
+
+# Foreign key columns from task_instance
+ti_fk_cols = ["dag_id", "task_id", "run_id", "map_index"]
+
+# Foreign key constraints from other tables to task_instance
+ti_fk_constrains = [
+ {"table": "rendered_task_instance_fields", "fk": "rtif_ti_fkey"},
+ {"table": "task_fail", "fk": "task_fail_ti_fkey"},
+ {"table": "task_instance_history", "fk": "task_instance_history_ti_fkey"},
+ {"table": "task_instance_note", "fk": "task_instance_note_ti_fkey"},
+ {"table": "task_map", "fk": "task_map_task_instance_fkey"},
+ {"table": "task_reschedule", "fk": "task_reschedule_ti_fkey"},
+ {"table": "xcom", "fk": "xcom_task_instance_fkey"},
+]
+
+
+def _get_type_id_column(dialect_name: str) -> sa.types.TypeEngine:
+ # For PostgreSQL, use the UUID type directly as it is more efficient
+ if dialect_name == "postgresql":
+ return postgresql.UUID(as_uuid=False)
+ # For other databases, use String(36) to match UUID format
+ else:
+ return sa.String(36)
+
+
+def upgrade():
+ """Add UUID primary key to task instance table."""
+ conn = op.get_bind()
+ dialect_name = conn.dialect.name
+
+ op.add_column("task_instance", sa.Column("id", _get_type_id_column(dialect_name), nullable=True))
+
+ if dialect_name == "postgresql":
+ op.execute(pg_uuid7_fn)
+
+ # TODO: Add batching to handle updates in smaller chunks for large tables to avoid locking
+ # Migrate existing rows with UUID v7 using a timestamp-based generation
+ op.execute(
+ "UPDATE task_instance SET id = uuid_generate_v7(coalesce(queued_dttm, start_date, clock_timestamp()))"
+ )
+
+ op.execute(pg_uuid7_fn_drop)
+
+ # Drop existing primary key constraint to task_instance table
+ op.execute("ALTER TABLE IF EXISTS task_instance DROP CONSTRAINT task_instance_pkey CASCADE")
+
+ elif dialect_name == "mysql":
+ op.execute(mysql_uuid7_fn)
+
+ # Migrate existing rows with UUID v7
+ op.execute("""
+ UPDATE task_instance
+ SET id = uuid_generate_v7(coalesce(queued_dttm, start_date, NOW(3)))
+ WHERE id IS NULL
+ """)
+
+ op.execute(mysql_uuid7_fn_drop)
Review Comment:
So these functions are used exclusively in the migrations? Might be worth adding a comment to make that clear.
##########
tests/jobs/test_local_task_job.py:
##########
@@ -902,7 +902,9 @@ def test_fast_follow(
ti_by_task_id = {}
with create_session() as session:
for task_id in init_state:
- ti = TaskInstance(dag.get_task(task_id), run_id=dag_run.run_id, state=init_state[task_id])
+ ti = TaskInstance(dag.get_task(task_id), run_id=dag_run.run_id)
+ ti.refresh_from_db()
Review Comment:
Wondering if the `refresh_from_db` could be moved to the `TaskInstance` constructor somehow as an optional operation?
##########
airflow/models/taskinstance.py:
##########
@@ -1813,10 +1817,16 @@ class TaskInstance(Base, LoggingMixin):
"""
__tablename__ = "task_instance"
- task_id = Column(StringID(), primary_key=True, nullable=False)
- dag_id = Column(StringID(), primary_key=True, nullable=False)
- run_id = Column(StringID(), primary_key=True, nullable=False)
- map_index = Column(Integer, primary_key=True, nullable=False, server_default=text("-1"))
+ id = Column(
+ String(36).with_variant(postgresql.UUID(as_uuid=False), "postgresql"),
+ primary_key=True,
+ default=lambda: str(uuid6.uuid7()),
Review Comment:
Declare as a top level function or classmethod so it can be used here and below (and easily swapped out for a different uuid7 implementation if needed)?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #43243:
URL: https://github.com/apache/airflow/pull/43243#issuecomment-2441404737
BTW. One of the ways it could be helped - we **could** potentially generate the unique id by hashing the remaining fields - with appropriate hashing algorithm, we could have very low probability of collision (and maybe we could implement a mechanism to detect collisions and implement handling of those collision in similar ways they are handled in hash maps) - that would be another way how we could approach it (and there the merge code coudl be simpler and still use session.merge()
But this one also comes with it's own set of difficulties and collision handling is likely going to be complex. But I thought it's worth menioning it here as an option.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil merged PR #43243:
URL: https://github.com/apache/airflow/pull/43243
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818239964
##########
airflow/api_connexion/endpoints/task_instance_endpoint.py:
##########
@@ -533,9 +533,11 @@ def post_set_task_instances_state(*, dag_id: str, session: Session = NEW_SESSION
detail=f"Task instance not found for task {task_id!r} on execution_date {execution_date}"
)
- if run_id and not session.get(
- TI, {"task_id": task_id, "dag_id": dag_id, "run_id": run_id, "map_index": -1}
- ):
+ select_stmt = select(TI).where(
+ TI.dag_id == dag_id, TI.task_id == task_id, TI.run_id == run_id, TI.map_index == -1
+ )
+
+ if run_id and not session.scalars(select_stmt).one_or_none():
Review Comment:
The reason for this and other changes in this file are explained in [**Issues identified Point (3)**](https://github.com/apache/airflow/pull/43243#issuecomment-2439774584)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make Task Instance primary key be a UUID [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #43243:
URL: https://github.com/apache/airflow/pull/43243#issuecomment-2428855292
A neat trick to avoid the need for pgcrypto extension: https://brandur.org/fragments/secure-bytes-without-pgcrypto
Though `uuid_send()` fn it uses doesn't exist on Pg10, it only comes from Pg11 onwards. However I do think it's reasonable to say 12 or 13 is the min required Pg version for Airflow 3 (Postgres only supports v13 onwards)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818240587
##########
tests/api_connexion/endpoints/test_task_instance_endpoint.py:
##########
@@ -1983,15 +1983,15 @@ def test_should_call_mocked_api(self, mock_set_task_instance_state, session):
self.create_task_instances(session)
NEW_STATE = "failed"
- mock_set_task_instance_state.return_value = session.get(
- TaskInstance,
- {
- "task_id": "print_the_context",
- "dag_id": "example_python_operator",
- "run_id": "TEST_DAG_RUN_ID",
- "map_index": -1,
- },
- )
+ mock_set_task_instance_state.return_value = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == "example_python_operator",
+ TaskInstance.task_id == "print_the_context",
+ TaskInstance.run_id == "TEST_DAG_RUN_ID",
+ TaskInstance.map_index == -1,
+ )
+ ).one_or_none()
+
Review Comment:
The reason for this and other changes in this file is explained in [**Issues identified Point (2)**](https://github.com/apache/airflow/pull/43243#issuecomment-2439774584)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate `TaskInstance` to UUID v7 primary key [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #43243:
URL: https://github.com/apache/airflow/pull/43243#discussion_r1818240857
##########
tests/models/test_cleartasks.py:
##########
@@ -597,7 +597,7 @@ def test_dag_clear(self, dag_maker):
ti0, ti1 = sorted(dr.task_instances, key=lambda ti: ti.task_id)
ti0.refresh_from_task(task0)
ti1.refresh_from_task(task1)
- session.get(TaskInstance, ti0.key.primary).try_number += 1
+ session.get(TaskInstance, ti0.id).try_number += 1
Review Comment:
The reason for this is explained in [**Issues identified Point (2)**](https://github.com/apache/airflow/pull/43243#issuecomment-2439774584)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org