From cd8d938356ac3d8908903978dd769ed808a52d5a Mon Sep 17 00:00:00 2001 From: hgbdev Date: Sat, 7 Dec 2024 01:10:47 +0700 Subject: [PATCH 1/4] fix(datasets): add created_by to ApiToken for get datasets from api correct current user (#11331) --- api/controllers/console/datasets/datasets.py | 1 + api/controllers/service_api/wraps.py | 23 +++++++++++--- ...42bf0de698_add_created_by_to_api_tokens.py | 31 +++++++++++++++++++ api/models/model.py | 1 + 4 files changed, 52 insertions(+), 4 deletions(-) create mode 100644 api/migrations/versions/2024_12_06_1324-5f42bf0de698_add_created_by_to_api_tokens.py diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index 95d4013e3a8f27..eb6a84e46e1a91 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/api/controllers/console/datasets/datasets.py @@ -566,6 +566,7 @@ def post(self): api_token.tenant_id = current_user.current_tenant_id api_token.token = key api_token.type = self.resource_type + api_token.created_by = current_user.id db.session.add(api_token) db.session.commit() return api_token, 200 diff --git a/api/controllers/service_api/wraps.py b/api/controllers/service_api/wraps.py index 2128c4c53f9909..62c8788ca603b3 100644 --- a/api/controllers/service_api/wraps.py +++ b/api/controllers/service_api/wraps.py @@ -140,14 +140,29 @@ def decorator(view): @wraps(view) def decorated(*args, **kwargs): api_token = validate_and_get_api_token("dataset") - tenant_account_join = ( + + # Build base query + query = ( db.session.query(Tenant, TenantAccountJoin) .filter(Tenant.id == api_token.tenant_id) .filter(TenantAccountJoin.tenant_id == Tenant.id) - .filter(TenantAccountJoin.role.in_(["owner"])) .filter(Tenant.status == TenantStatus.NORMAL) - .one_or_none() - ) # TODO: only owner information is required, so only one is returned. + ) + + if api_token.created_by: + # Only apply account_id filter if created_by exists + query = query.filter( + db.and_( + TenantAccountJoin.role.in_(["owner", "admin"]), + TenantAccountJoin.account_id == api_token.created_by, + ) + ) + else: + query = query.filter(TenantAccountJoin.role.in_(["owner"])) + + tenant_account_join = query.one_or_none() + # TODO: only owner information is required, so only one is returned. + if tenant_account_join: tenant, ta = tenant_account_join account = Account.query.filter_by(id=ta.account_id).first() diff --git a/api/migrations/versions/2024_12_06_1324-5f42bf0de698_add_created_by_to_api_tokens.py b/api/migrations/versions/2024_12_06_1324-5f42bf0de698_add_created_by_to_api_tokens.py new file mode 100644 index 00000000000000..cc2d77859be4ab --- /dev/null +++ b/api/migrations/versions/2024_12_06_1324-5f42bf0de698_add_created_by_to_api_tokens.py @@ -0,0 +1,31 @@ +"""add created_by to api_tokens + +Revision ID: 5f42bf0de698 +Revises: 01d6889832f7 +Create Date: 2024-12-06 13:24:28.701384 + +""" +from alembic import op +import models as models +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '5f42bf0de698' +down_revision = '01d6889832f7' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('api_tokens', schema=None) as batch_op: + batch_op.add_column(sa.Column('created_by', models.types.StringUUID(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('api_tokens', schema=None) as batch_op: + batch_op.drop_column('created_by') + # ### end Alembic commands ### diff --git a/api/models/model.py b/api/models/model.py index 03b8e0bea553aa..6f403ab23aa10b 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -1340,6 +1340,7 @@ class ApiToken(db.Model): type = db.Column(db.String(16), nullable=False) token = db.Column(db.String(255), nullable=False) last_used_at = db.Column(db.DateTime, nullable=True) + created_by = db.Column(StringUUID, nullable=True) created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @staticmethod From c9c9ca76fd973199a7c15974f7dac6c80acc45a4 Mon Sep 17 00:00:00 2001 From: hgbdev Date: Sat, 7 Dec 2024 01:38:42 +0700 Subject: [PATCH 2/4] fix(datasets): only return ApiKey of user who created (#11331) --- api/controllers/console/datasets/datasets.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index eb6a84e46e1a91..38a0e93104c2b1 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/api/controllers/console/datasets/datasets.py @@ -537,6 +537,9 @@ def get(self): .filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id) .all() ) + + keys = [key for key in keys if key.created_by == current_user.id or key.created_by is None] + return {"items": keys} @setup_required @@ -548,12 +551,14 @@ def post(self): if not current_user.is_admin_or_owner: raise Forbidden() - current_key_count = ( + keys = ( db.session.query(ApiToken) .filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id) - .count() + .all() ) + current_key_count = len([key for key in keys if key.created_by == current_user.id or key.created_by is None]) + if current_key_count >= self.max_keys: flask_restful.abort( 400, From 4a1777c93859f1e641de881529f2bea39df73477 Mon Sep 17 00:00:00 2001 From: hgbdev Date: Sun, 8 Dec 2024 02:41:23 +0700 Subject: [PATCH 3/4] chore: add created_by index to api_tokens table --- ..._1936-8db4e7683504_add_created_by_to_api_tokens.py} | 10 +++++++--- api/models/model.py | 1 + 2 files changed, 8 insertions(+), 3 deletions(-) rename api/migrations/versions/{2024_12_06_1324-5f42bf0de698_add_created_by_to_api_tokens.py => 2024_12_07_1936-8db4e7683504_add_created_by_to_api_tokens.py} (76%) diff --git a/api/migrations/versions/2024_12_06_1324-5f42bf0de698_add_created_by_to_api_tokens.py b/api/migrations/versions/2024_12_07_1936-8db4e7683504_add_created_by_to_api_tokens.py similarity index 76% rename from api/migrations/versions/2024_12_06_1324-5f42bf0de698_add_created_by_to_api_tokens.py rename to api/migrations/versions/2024_12_07_1936-8db4e7683504_add_created_by_to_api_tokens.py index cc2d77859be4ab..a7c833b8ff2ded 100644 --- a/api/migrations/versions/2024_12_06_1324-5f42bf0de698_add_created_by_to_api_tokens.py +++ b/api/migrations/versions/2024_12_07_1936-8db4e7683504_add_created_by_to_api_tokens.py @@ -1,8 +1,8 @@ """add created_by to api_tokens -Revision ID: 5f42bf0de698 +Revision ID: 8db4e7683504 Revises: 01d6889832f7 -Create Date: 2024-12-06 13:24:28.701384 +Create Date: 2024-12-07 19:36:49.632151 """ from alembic import op @@ -11,7 +11,7 @@ # revision identifiers, used by Alembic. -revision = '5f42bf0de698' +revision = '8db4e7683504' down_revision = '01d6889832f7' branch_labels = None depends_on = None @@ -21,11 +21,15 @@ def upgrade(): # ### commands auto generated by Alembic - please adjust! ### with op.batch_alter_table('api_tokens', schema=None) as batch_op: batch_op.add_column(sa.Column('created_by', models.types.StringUUID(), nullable=True)) + batch_op.create_index('api_token_created_by_idx', ['created_by'], unique=False) + # ### end Alembic commands ### def downgrade(): # ### commands auto generated by Alembic - please adjust! ### with op.batch_alter_table('api_tokens', schema=None) as batch_op: + batch_op.drop_index('api_token_created_by_idx') batch_op.drop_column('created_by') + # ### end Alembic commands ### diff --git a/api/models/model.py b/api/models/model.py index 6f403ab23aa10b..8dfa0dc162a759 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -1332,6 +1332,7 @@ class ApiToken(db.Model): db.Index("api_token_app_id_type_idx", "app_id", "type"), db.Index("api_token_token_idx", "token", "type"), db.Index("api_token_tenant_idx", "tenant_id", "type"), + db.Index("api_token_created_by_idx", "created_by"), ) id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) From 4c20b9205ebd0ccec2df54331fc95395189cac18 Mon Sep 17 00:00:00 2001 From: hgbdev Date: Sun, 8 Dec 2024 02:51:50 +0700 Subject: [PATCH 4/4] chore: add clarifying comments for dataset API key filtering based on created_by field --- api/controllers/console/datasets/datasets.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index 38a0e93104c2b1..30ea1f1de4dbce 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/api/controllers/console/datasets/datasets.py @@ -538,6 +538,7 @@ def get(self): .all() ) + # Only return keys created by current user or created_by none (for old keys do not have created_by) keys = [key for key in keys if key.created_by == current_user.id or key.created_by is None] return {"items": keys} @@ -557,6 +558,7 @@ def post(self): .all() ) + # Only count keys created by current user or created_by none (for old keys do not have created_by) current_key_count = len([key for key in keys if key.created_by == current_user.id or key.created_by is None]) if current_key_count >= self.max_keys: