Skip to content

Commit 41899d8

Browse files
jake11-ohoclaude
jake11-oho
and
claude
committed
feat(datasets): add DatasetMetadataClient for pure DB-backed metadata management
Introduce a pure DB-backed metadata client decoupled from the server package. ORM models (Dataset, Instance, Image, Permission, AuditEvent) are defined in rock/sdk/envhub/datasets/database.py with their own Base, eliminating the import chain through rock.envhub that pulled in fastapi, nacos, redis, and ray. Key changes: - DatasetMetadataClient as user-facing SDK entry point - DbDatasetRegistry with connection pooling, engine dedup, batch upsert - pool_pre_ping parameter, create_all dedup per db_url - Split table with task_count, tags on Instance, timestamp fields - Sorting support (default: org+name ASC) with SortField/SortOrder enums - DatasetInfo.splits returns list[SplitInfo] with full metadata - source field on Dataset for tracking data origin - 63 unit tests covering all CRUD operations, sorting, and edge cases fixes #1170 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 8ae2a22 commit 41899d8

8 files changed

Lines changed: 3254 additions & 2 deletions

File tree

‎docs/proposals/dataset-metadata-client.md‎

Lines changed: 934 additions & 0 deletions
Large diffs are not rendered by default.

‎rock/envhub/database/models.py‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
from __future__ import annotations
2+
3+
from sqlalchemy import (
4+
JSON,
5+
BigInteger,
6+
Column,
7+
DateTime,
8+
ForeignKey,
9+
Index,
10+
Integer,
11+
String,
12+
Text,
13+
UniqueConstraint,
14+
func,
15+
)
16+
from sqlalchemy.orm import relationship
17+
18+
from rock.envhub.database.base import Base
19+
20+
21+
class Dataset(Base):
22+
__tablename__ = "datasets"
23+
24+
id = Column(Integer, primary_key=True, autoincrement=True)
25+
org = Column(String(255), nullable=False, index=True)
26+
name = Column(String(255), nullable=False)
27+
description = Column(Text, default="")
28+
tags = Column(JSON, default=list)
29+
owner = Column(String(255), default="")
30+
homepage = Column(String(512), nullable=True)
31+
repo = Column(String(512), nullable=True)
32+
paper = Column(String(512), nullable=True)
33+
leaderboard = Column(String(512), nullable=True)
34+
logo_url = Column(String(512), nullable=True)
35+
os = Column(String(64), nullable=True)
36+
version = Column(String(64), nullable=True)
37+
task_counts = Column(JSON, nullable=True, default=dict)
38+
created_at = Column(DateTime, server_default=func.now())
39+
updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
40+
41+
instances = relationship("Instance", back_populates="dataset", cascade="all, delete-orphan")
42+
permissions = relationship("DatasetPermission", back_populates="dataset", cascade="all, delete-orphan")
43+
44+
__table_args__ = (UniqueConstraint("org", "name", name="uq_dataset_org_name"),)
45+
46+
@property
47+
def full_name(self) -> str:
48+
return f"{self.org}/{self.name}"
49+
50+
def __repr__(self):
51+
return f"<Dataset(id={self.id}, name='{self.full_name}')>"
52+
53+
54+
class Instance(Base):
55+
__tablename__ = "instances"
56+
57+
id = Column(Integer, primary_key=True, autoincrement=True)
58+
dataset_id = Column(Integer, ForeignKey("datasets.id", ondelete="CASCADE"), nullable=False)
59+
split = Column(String(255), nullable=False)
60+
name = Column(String(255), nullable=False)
61+
description = Column(Text, default="")
62+
type = Column(String(16), default="directory")
63+
size = Column(BigInteger, nullable=True)
64+
file_count = Column(Integer, nullable=True)
65+
etag = Column(String(255), nullable=True)
66+
format = Column(String(64), nullable=True, index=True)
67+
repo = Column(String(512), nullable=True)
68+
language = Column(String(64), nullable=True, index=True)
69+
difficulty = Column(String(64), nullable=True)
70+
base_commit = Column(String(64), nullable=True)
71+
image_uris = Column(JSON, nullable=True)
72+
raw = Column(Text, nullable=True)
73+
source_revision = Column(String(128), nullable=True)
74+
imported_from = Column(String(512), nullable=True)
75+
created_by = Column(String(255), nullable=True)
76+
created_at = Column(DateTime, server_default=func.now())
77+
updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
78+
79+
dataset = relationship("Dataset", back_populates="instances")
80+
81+
__table_args__ = (
82+
UniqueConstraint("dataset_id", "split", "name", name="uq_instance_dataset_split_name"),
83+
Index("ix_instance_dataset_split", "dataset_id", "split"),
84+
)
85+
86+
def __repr__(self):
87+
return f"<Instance(id={self.id}, name='{self.name}', split='{self.split}')>"
88+
89+
90+
class Image(Base):
91+
__tablename__ = "images"
92+
93+
source_image_uri = Column(String(512), primary_key=True)
94+
image_uri_sg = Column(String(512), nullable=True)
95+
image_uri_sh = Column(String(512), nullable=True)
96+
image_hash = Column(String(71), nullable=True)
97+
status = Column(String(32), nullable=False, default="pending", index=True)
98+
last_error = Column(Text, nullable=True)
99+
last_job_id = Column(String(64), nullable=True)
100+
created_by = Column(String(255), nullable=False, default="system")
101+
created_at = Column(DateTime, server_default=func.now())
102+
updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
103+
104+
def __repr__(self):
105+
return f"<Image(source_image_uri='{self.source_image_uri}', status='{self.status}')>"
106+
107+
108+
class DatasetPermission(Base):
109+
__tablename__ = "dataset_permissions"
110+
111+
id = Column(Integer, primary_key=True, autoincrement=True)
112+
dataset_id = Column(Integer, ForeignKey("datasets.id", ondelete="CASCADE"), nullable=False)
113+
user_id = Column(String(255), nullable=False)
114+
role = Column(String(32), nullable=False, default="viewer")
115+
granted_by = Column(String(255), nullable=True)
116+
created_at = Column(DateTime, server_default=func.now())
117+
updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
118+
119+
dataset = relationship("Dataset", back_populates="permissions")
120+
121+
__table_args__ = (
122+
UniqueConstraint("dataset_id", "user_id", name="uq_permission_dataset_user"),
123+
Index("ix_permission_user", "user_id"),
124+
)
125+
126+
def __repr__(self):
127+
return f"<DatasetPermission(dataset_id={self.dataset_id}, user_id='{self.user_id}', role='{self.role}')>"
128+
129+
130+
class AuditEvent(Base):
131+
__tablename__ = "audit_events"
132+
133+
id = Column(Integer, primary_key=True, autoincrement=True)
134+
target_type = Column(String(32), nullable=False, index=True)
135+
target_id = Column(String(512), nullable=False)
136+
event_type = Column(String(64), nullable=False, index=True)
137+
operator = Column(String(255), nullable=False)
138+
changes = Column(JSON, nullable=True)
139+
created_at = Column(DateTime, server_default=func.now())
140+
141+
__table_args__ = (Index("ix_audit_target", "target_type", "target_id"),)
142+
143+
def __repr__(self):
144+
return f"<AuditEvent(id={self.id}, target_type='{self.target_type}', event_type='{self.event_type}')>"

‎rock/sdk/envhub/datasets/__init__.py‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,31 @@
11
from rock.sdk.envhub.datasets.client import DatasetClient
2-
from rock.sdk.envhub.datasets.models import DatasetSpec, UploadResult
2+
from rock.sdk.envhub.datasets.metadata_client import DatasetMetadataClient
3+
from rock.sdk.envhub.datasets.models import (
4+
AuditEventInfo,
5+
DatasetInfo,
6+
DatasetSpec,
7+
ImageInfo,
8+
PageResult,
9+
PermissionInfo,
10+
SortField,
11+
SortOrder,
12+
SplitInfo,
13+
TaskEntry,
14+
UploadResult,
15+
)
316

4-
__all__ = ["DatasetClient", "DatasetSpec", "UploadResult"]
17+
__all__ = [
18+
"DatasetClient",
19+
"DatasetMetadataClient",
20+
"AuditEventInfo",
21+
"DatasetInfo",
22+
"DatasetSpec",
23+
"ImageInfo",
24+
"PageResult",
25+
"PermissionInfo",
26+
"SortField",
27+
"SortOrder",
28+
"SplitInfo",
29+
"TaskEntry",
30+
"UploadResult",
31+
]

‎rock/sdk/envhub/datasets/database.py‎

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
from __future__ import annotations
2+
3+
from sqlalchemy import (
4+
JSON,
5+
BigInteger,
6+
Column,
7+
DateTime,
8+
ForeignKey,
9+
Index,
10+
Integer,
11+
String,
12+
Text,
13+
UniqueConstraint,
14+
func,
15+
)
16+
from sqlalchemy.orm import DeclarativeBase, relationship
17+
18+
19+
class Base(DeclarativeBase):
20+
pass
21+
22+
23+
class Dataset(Base):
24+
__tablename__ = "datasets"
25+
26+
id = Column(Integer, primary_key=True, autoincrement=True)
27+
org = Column(String(255), nullable=False, index=True)
28+
name = Column(String(255), nullable=False)
29+
source = Column(String(128), default="", server_default="")
30+
description = Column(Text, default="")
31+
tags = Column(JSON, default=list)
32+
owner = Column(String(255), default="")
33+
homepage = Column(String(512), nullable=True)
34+
repo = Column(String(512), nullable=True)
35+
paper = Column(String(512), nullable=True)
36+
leaderboard = Column(String(512), nullable=True)
37+
logo_url = Column(String(512), nullable=True)
38+
os = Column(String(64), nullable=True)
39+
version = Column(String(64), nullable=True)
40+
created_at = Column(DateTime, server_default=func.now())
41+
updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
42+
43+
instances = relationship("Instance", back_populates="dataset", cascade="all, delete-orphan")
44+
permissions = relationship("DatasetPermission", back_populates="dataset", cascade="all, delete-orphan")
45+
splits_rel = relationship("Split", back_populates="dataset", cascade="all, delete-orphan")
46+
47+
__table_args__ = (UniqueConstraint("org", "name", name="uq_dataset_org_name"),)
48+
49+
@property
50+
def full_name(self) -> str:
51+
return f"{self.org}/{self.name}"
52+
53+
def __repr__(self):
54+
return f"<Dataset(id={self.id}, name='{self.full_name}')>"
55+
56+
57+
class Instance(Base):
58+
__tablename__ = "instances"
59+
60+
id = Column(Integer, primary_key=True, autoincrement=True)
61+
dataset_id = Column(Integer, ForeignKey("datasets.id", ondelete="CASCADE"), nullable=False)
62+
split = Column(String(255), nullable=False)
63+
name = Column(String(255), nullable=False)
64+
description = Column(Text, default="")
65+
type = Column(String(16), default="directory")
66+
size = Column(BigInteger, nullable=True)
67+
file_count = Column(Integer, nullable=True)
68+
etag = Column(String(255), nullable=True)
69+
format = Column(String(64), nullable=True, index=True)
70+
repo = Column(String(512), nullable=True)
71+
language = Column(String(64), nullable=True, index=True)
72+
difficulty = Column(String(64), nullable=True)
73+
base_commit = Column(String(64), nullable=True)
74+
image_uris = Column(JSON, nullable=True)
75+
tags = Column(JSON, default=list)
76+
raw = Column(Text, nullable=True)
77+
source_revision = Column(String(128), nullable=True)
78+
imported_from = Column(String(512), nullable=True)
79+
created_by = Column(String(255), nullable=True)
80+
created_at = Column(DateTime, server_default=func.now())
81+
updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
82+
83+
dataset = relationship("Dataset", back_populates="instances")
84+
85+
__table_args__ = (
86+
UniqueConstraint("dataset_id", "split", "name", name="uq_instance_dataset_split_name"),
87+
Index("ix_instance_dataset_split", "dataset_id", "split"),
88+
)
89+
90+
def __repr__(self):
91+
return f"<Instance(id={self.id}, name='{self.name}', split='{self.split}')>"
92+
93+
94+
class Split(Base):
95+
__tablename__ = "splits"
96+
97+
id = Column(Integer, primary_key=True, autoincrement=True)
98+
dataset_id = Column(Integer, ForeignKey("datasets.id", ondelete="CASCADE"), nullable=False)
99+
name = Column(String(255), nullable=False)
100+
task_count = Column(Integer, nullable=False, default=0)
101+
created_by = Column(String(255), nullable=True)
102+
created_at = Column(DateTime, server_default=func.now())
103+
updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
104+
105+
dataset = relationship("Dataset", back_populates="splits_rel")
106+
107+
__table_args__ = (
108+
UniqueConstraint("dataset_id", "name", name="uq_split_dataset_name"),
109+
Index("ix_split_dataset_name", "dataset_id", "name"),
110+
)
111+
112+
def __repr__(self):
113+
return f"<Split(id={self.id}, dataset_id={self.dataset_id}, name='{self.name}', task_count={self.task_count})>"
114+
115+
116+
class Image(Base):
117+
__tablename__ = "images"
118+
119+
source_image_uri = Column(String(512), primary_key=True)
120+
image_uri_sg = Column(String(512), nullable=True)
121+
image_uri_sh = Column(String(512), nullable=True)
122+
image_hash = Column(String(71), nullable=True)
123+
status = Column(String(32), nullable=False, default="pending", index=True)
124+
last_error = Column(Text, nullable=True)
125+
last_job_id = Column(String(64), nullable=True)
126+
created_by = Column(String(255), nullable=False, default="system")
127+
created_at = Column(DateTime, server_default=func.now())
128+
updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
129+
130+
def __repr__(self):
131+
return f"<Image(source_image_uri='{self.source_image_uri}', status='{self.status}')>"
132+
133+
134+
class DatasetPermission(Base):
135+
__tablename__ = "dataset_permissions"
136+
137+
id = Column(Integer, primary_key=True, autoincrement=True)
138+
dataset_id = Column(Integer, ForeignKey("datasets.id", ondelete="CASCADE"), nullable=False)
139+
user_id = Column(String(255), nullable=False)
140+
role = Column(String(32), nullable=False, default="viewer")
141+
granted_by = Column(String(255), nullable=True)
142+
created_at = Column(DateTime, server_default=func.now())
143+
updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
144+
145+
dataset = relationship("Dataset", back_populates="permissions")
146+
147+
__table_args__ = (
148+
UniqueConstraint("dataset_id", "user_id", name="uq_permission_dataset_user"),
149+
Index("ix_permission_user", "user_id"),
150+
)
151+
152+
def __repr__(self):
153+
return f"<DatasetPermission(dataset_id={self.dataset_id}, user_id='{self.user_id}', role='{self.role}')>"
154+
155+
156+
class AuditEvent(Base):
157+
__tablename__ = "audit_events"
158+
159+
id = Column(Integer, primary_key=True, autoincrement=True)
160+
target_type = Column(String(32), nullable=False, index=True)
161+
target_id = Column(String(512), nullable=False)
162+
event_type = Column(String(64), nullable=False, index=True)
163+
operator = Column(String(255), nullable=False)
164+
changes = Column(JSON, nullable=True)
165+
created_at = Column(DateTime, server_default=func.now())
166+
167+
__table_args__ = (Index("ix_audit_target", "target_type", "target_id"),)
168+
169+
def __repr__(self):
170+
return f"<AuditEvent(id={self.id}, target_type='{self.target_type}', event_type='{self.event_type}')>"

0 commit comments

Comments
 (0)