#!/usr/bin/env node
import { Datastore } from "@google-cloud/datastore";
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import {
CallToolRequestSchema,
ListToolsRequestSchema,
Tool,
} from "@modelcontextprotocol/sdk/types.js";
// Initialize Datastore client
const datastore = new Datastore({
projectId: process.env.GOOGLE_CLOUD_PROJECT,
keyFilename: process.env.GOOGLE_APPLICATION_CREDENTIALS,
});
// Define available tools
const tools: Tool[] = [
{
name: "datastore_get",
description: "Get an entity by kind and key from Datastore",
inputSchema: {
type: "object",
properties: {
kind: {
type: "string",
description: "The kind (type) of the entity",
},
keyId: {
type: ["string", "number"],
description: "The ID of the key (can be string name or numeric ID)",
},
namespace: {
type: "string",
description: "Optional namespace for the entity",
},
ancestors: {
type: "array",
description: "Optional array of ancestor keys [kind, id/name, kind, id/name, ...]",
items: {
type: ["string", "number"]
}
}
},
required: ["kind", "keyId"],
},
},
{
name: "datastore_insert",
description: "Insert a new entity into Datastore",
inputSchema: {
type: "object",
properties: {
kind: {
type: "string",
description: "The kind (type) of the entity",
},
data: {
type: "object",
description: "The entity data to insert",
},
keyId: {
type: ["string", "number"],
description: "Optional key ID (if not provided, Datastore will auto-generate)",
},
namespace: {
type: "string",
description: "Optional namespace for the entity",
},
excludeFromIndexes: {
type: "array",
description: "Optional array of property names to exclude from indexes",
items: {
type: "string"
}
}
},
required: ["kind", "data"],
},
},
{
name: "datastore_update",
description: "Update an existing entity in Datastore",
inputSchema: {
type: "object",
properties: {
kind: {
type: "string",
description: "The kind (type) of the entity",
},
keyId: {
type: ["string", "number"],
description: "The ID of the key",
},
data: {
type: "object",
description: "The updated entity data",
},
namespace: {
type: "string",
description: "Optional namespace for the entity",
},
excludeFromIndexes: {
type: "array",
description: "Optional array of property names to exclude from indexes",
items: {
type: "string"
}
}
},
required: ["kind", "keyId", "data"],
},
},
{
name: "datastore_upsert",
description: "Insert or update an entity in Datastore (upsert operation)",
inputSchema: {
type: "object",
properties: {
kind: {
type: "string",
description: "The kind (type) of the entity",
},
keyId: {
type: ["string", "number"],
description: "The ID of the key",
},
data: {
type: "object",
description: "The entity data",
},
namespace: {
type: "string",
description: "Optional namespace for the entity",
},
excludeFromIndexes: {
type: "array",
description: "Optional array of property names to exclude from indexes",
items: {
type: "string"
}
}
},
required: ["kind", "keyId", "data"],
},
},
{
name: "datastore_delete",
description: "Delete an entity from Datastore",
inputSchema: {
type: "object",
properties: {
kind: {
type: "string",
description: "The kind (type) of the entity",
},
keyId: {
type: ["string", "number"],
description: "The ID of the key",
},
namespace: {
type: "string",
description: "Optional namespace for the entity",
}
},
required: ["kind", "keyId"],
},
},
{
name: "datastore_query",
description: "Query entities from Datastore with optional filters, ordering, and pagination",
inputSchema: {
type: "object",
properties: {
kind: {
type: "string",
description: "The kind (type) of entities to query",
},
filters: {
type: "array",
description: "Optional array of filter conditions",
items: {
type: "object",
properties: {
property: {
type: "string",
description: "The property name to filter on",
},
operator: {
type: "string",
description: "Comparison operator: =, <, >, <=, >=, HAS_ANCESTOR",
enum: ["=", "<", ">", "<=", ">=", "HAS_ANCESTOR"],
},
value: {
description: "The value to compare against",
},
},
required: ["property", "operator", "value"],
},
},
orderBy: {
type: "array",
description: "Optional array of properties to order by",
items: {
type: "object",
properties: {
property: {
type: "string",
description: "Property name to order by",
},
descending: {
type: "boolean",
description: "Whether to order descending (default: false)",
},
},
required: ["property"],
},
},
limit: {
type: "number",
description: "Maximum number of results to return",
},
offset: {
type: "number",
description: "Number of results to skip",
},
namespace: {
type: "string",
description: "Optional namespace to query from",
},
select: {
type: "array",
description: "Optional array of property names to return (projection query)",
items: {
type: "string"
}
}
},
required: ["kind"],
},
},
{
name: "datastore_runAggregationQuery",
description: "Run an aggregation query (count, sum, avg) on Datastore entities",
inputSchema: {
type: "object",
properties: {
kind: {
type: "string",
description: "The kind (type) of entities to query",
},
aggregations: {
type: "array",
description: "Array of aggregations to perform",
items: {
type: "object",
properties: {
type: {
type: "string",
description: "Type of aggregation: count, sum, avg",
enum: ["count", "sum", "avg"],
},
property: {
type: "string",
description: "Property to aggregate (required for sum and avg)",
},
alias: {
type: "string",
description: "Optional alias for the aggregation result",
},
},
required: ["type"],
},
},
filters: {
type: "array",
description: "Optional array of filter conditions",
items: {
type: "object",
properties: {
property: {
type: "string",
},
operator: {
type: "string",
enum: ["=", "<", ">", "<=", ">="],
},
value: {},
},
required: ["property", "operator", "value"],
},
},
namespace: {
type: "string",
description: "Optional namespace to query from",
}
},
required: ["kind", "aggregations"],
},
},
{
name: "datastore_allocateIds",
description: "Allocate IDs for incomplete keys",
inputSchema: {
type: "object",
properties: {
kind: {
type: "string",
description: "The kind (type) for the keys",
},
count: {
type: "number",
description: "Number of IDs to allocate",
},
namespace: {
type: "string",
description: "Optional namespace",
}
},
required: ["kind", "count"],
},
},
{
name: "datastore_createKey",
description: "Create a complete or incomplete key",
inputSchema: {
type: "object",
properties: {
kind: {
type: "string",
description: "The kind (type) for the key",
},
keyId: {
type: ["string", "number"],
description: "Optional key ID (omit for incomplete key)",
},
namespace: {
type: "string",
description: "Optional namespace",
},
path: {
type: "array",
description: "Optional array for creating hierarchical keys [kind, id/name, kind, id/name, ...]",
items: {
type: ["string", "number"]
}
}
},
required: ["kind"],
},
},
{
name: "datastore_runInTransaction",
description: "Execute multiple operations within a transaction",
inputSchema: {
type: "object",
properties: {
operations: {
type: "array",
description: "Array of operations to execute in the transaction",
items: {
type: "object",
properties: {
type: {
type: "string",
description: "Operation type: get, insert, update, upsert, delete",
enum: ["get", "insert", "update", "upsert", "delete"],
},
kind: {
type: "string",
description: "The kind (type) of the entity",
},
keyId: {
type: ["string", "number"],
description: "The key ID (not required for insert)",
},
data: {
type: "object",
description: "Entity data (for insert, update, upsert)",
},
namespace: {
type: "string",
description: "Optional namespace",
}
},
required: ["type", "kind"],
},
},
},
required: ["operations"],
},
},
{
name: "datastore_listKinds",
description: "List all entity kinds (types) in the Datastore by querying the __kind__ metadata",
inputSchema: {
type: "object",
properties: {
namespace: {
type: "string",
description: "Optional namespace to list kinds from",
}
},
},
},
];
// Create server instance
const server = new Server(
{
name: "mcp-datastore-server",
version: "1.0.0",
},
{
capabilities: {
tools: {},
},
}
);
// Helper function to create a key
function createKey(
kind: string,
keyId?: string | number,
namespace?: string,
ancestors?: (string | number)[]
) {
const options: any = {};
if (namespace) {
options.namespace = namespace;
}
let path: (string | number)[] = [];
// Add ancestors if provided
if (ancestors && ancestors.length > 0) {
path = [...ancestors];
}
// Add the main kind and keyId
path.push(kind);
if (keyId !== undefined) {
path.push(keyId);
}
if (path.length > 0) {
options.path = path;
}
return datastore.key(options);
}
// Helper function to serialize entity for response
function serializeEntity(entity: any) {
if (!entity) return null;
const result: any = {};
// Add key information
if (entity[datastore.KEY]) {
const key = entity[datastore.KEY];
result._key = {
kind: key.kind,
id: key.id,
name: key.name,
namespace: key.namespace,
path: key.path,
};
}
// Add entity properties
for (const [key, value] of Object.entries(entity)) {
if (key !== datastore.KEY.toString()) {
result[key] = value;
}
}
return result;
}
// Handle list tools request
server.setRequestHandler(ListToolsRequestSchema, async () => {
return { tools };
});
// Handle tool execution
server.setRequestHandler(CallToolRequestSchema, async (request) => {
const { name, arguments: args } = request.params;
try {
switch (name) {
case "datastore_get": {
const { kind, keyId, namespace, ancestors } = args as any;
const key = createKey(kind, keyId, namespace, ancestors);
const [entity] = await datastore.get(key);
return {
content: [
{
type: "text",
text: entity
? JSON.stringify(serializeEntity(entity), null, 2)
: "Entity not found",
},
],
};
}
case "datastore_insert": {
const { kind, data, keyId, namespace, excludeFromIndexes } = args as any;
const key = createKey(kind, keyId, namespace);
const entity: any = {
key,
data,
};
if (excludeFromIndexes && excludeFromIndexes.length > 0) {
entity.excludeFromIndexes = excludeFromIndexes;
}
await datastore.insert(entity);
return {
content: [
{
type: "text",
text: JSON.stringify({
success: true,
message: "Entity inserted successfully",
key: {
kind: key.kind,
id: key.id,
name: key.name,
},
}, null, 2),
},
],
};
}
case "datastore_update": {
const { kind, keyId, data, namespace, excludeFromIndexes } = args as any;
const key = createKey(kind, keyId, namespace);
const entity: any = {
key,
data,
};
if (excludeFromIndexes && excludeFromIndexes.length > 0) {
entity.excludeFromIndexes = excludeFromIndexes;
}
await datastore.update(entity);
return {
content: [
{
type: "text",
text: JSON.stringify({
success: true,
message: "Entity updated successfully",
}, null, 2),
},
],
};
}
case "datastore_upsert": {
const { kind, keyId, data, namespace, excludeFromIndexes } = args as any;
const key = createKey(kind, keyId, namespace);
const entity: any = {
key,
data,
};
if (excludeFromIndexes && excludeFromIndexes.length > 0) {
entity.excludeFromIndexes = excludeFromIndexes;
}
await datastore.upsert(entity);
return {
content: [
{
type: "text",
text: JSON.stringify({
success: true,
message: "Entity upserted successfully",
}, null, 2),
},
],
};
}
case "datastore_delete": {
const { kind, keyId, namespace } = args as any;
const key = createKey(kind, keyId, namespace);
await datastore.delete(key);
return {
content: [
{
type: "text",
text: JSON.stringify({
success: true,
message: "Entity deleted successfully",
}, null, 2),
},
],
};
}
case "datastore_query": {
const { kind, filters, orderBy, limit, offset, namespace, select } = args as any;
let query = datastore.createQuery(namespace, kind);
// Apply filters
if (filters && filters.length > 0) {
for (const filter of filters) {
query = query.filter(filter.property, filter.operator, filter.value);
}
}
// Apply ordering
if (orderBy && orderBy.length > 0) {
for (const order of orderBy) {
query = query.order(order.property, {
descending: order.descending || false,
});
}
}
// Apply limit
if (limit) {
query = query.limit(limit);
}
// Apply offset
if (offset) {
query = query.offset(offset);
}
// Apply select (projection)
if (select && select.length > 0) {
query = query.select(select);
}
const [entities] = await datastore.runQuery(query);
return {
content: [
{
type: "text",
text: JSON.stringify({
count: entities.length,
entities: entities.map(serializeEntity),
}, null, 2),
},
],
};
}
case "datastore_runAggregationQuery": {
const { kind, aggregations, filters, namespace } = args as any;
let query = datastore.createQuery(namespace, kind);
// Apply filters
if (filters && filters.length > 0) {
for (const filter of filters) {
query = query.filter(filter.property, filter.operator, filter.value);
}
}
// Create aggregation query
let aggQuery = datastore.createAggregationQuery(query);
// Apply aggregations
for (const agg of aggregations) {
const alias = agg.alias || `${agg.type}_${agg.property || 'result'}`;
switch (agg.type) {
case "count":
aggQuery = aggQuery.count(alias);
break;
case "sum":
if (!agg.property) {
throw new Error("Property is required for sum aggregation");
}
aggQuery = aggQuery.sum(agg.property, alias);
break;
case "avg":
if (!agg.property) {
throw new Error("Property is required for avg aggregation");
}
aggQuery = aggQuery.average(agg.property, alias);
break;
}
}
const [results] = await datastore.runAggregationQuery(aggQuery);
return {
content: [
{
type: "text",
text: JSON.stringify(results, null, 2),
},
],
};
}
case "datastore_allocateIds": {
const { kind, count, namespace } = args as any;
const incompleteKey = createKey(kind, undefined, namespace);
const [keys] = await datastore.allocateIds(incompleteKey, count);
return {
content: [
{
type: "text",
text: JSON.stringify({
allocated: keys.map((key: any) => ({
kind: key.kind,
id: key.id,
name: key.name,
path: key.path,
})),
}, null, 2),
},
],
};
}
case "datastore_createKey": {
const { kind, keyId, namespace, path } = args as any;
let key;
if (path && path.length > 0) {
// Create key with path
key = createKey(kind, keyId, namespace, path);
} else {
key = createKey(kind, keyId, namespace);
}
return {
content: [
{
type: "text",
text: JSON.stringify({
kind: key.kind,
id: key.id,
name: key.name,
namespace: key.namespace,
path: key.path,
}, null, 2),
},
],
};
}
case "datastore_runInTransaction": {
const { operations } = args as any;
const transaction = datastore.transaction();
await transaction.run();
try {
const results: any[] = [];
for (const op of operations) {
const { type, kind, keyId, data, namespace } = op;
const key = createKey(kind, keyId, namespace);
switch (type) {
case "get": {
const [entity] = await transaction.get(key);
results.push({
operation: "get",
entity: serializeEntity(entity),
});
break;
}
case "insert": {
const entity: any = { key, data };
transaction.insert(entity);
results.push({
operation: "insert",
key: { kind: key.kind, id: key.id, name: key.name },
});
break;
}
case "update": {
const entity: any = { key, data };
transaction.update(entity);
results.push({
operation: "update",
key: { kind: key.kind, id: key.id, name: key.name },
});
break;
}
case "upsert": {
const entity: any = { key, data };
transaction.upsert(entity);
results.push({
operation: "upsert",
key: { kind: key.kind, id: key.id, name: key.name },
});
break;
}
case "delete": {
transaction.delete(key);
results.push({
operation: "delete",
key: { kind: key.kind, id: key.id, name: key.name },
});
break;
}
}
}
await transaction.commit();
return {
content: [
{
type: "text",
text: JSON.stringify({
success: true,
message: "Transaction completed successfully",
results,
}, null, 2),
},
],
};
} catch (error) {
await transaction.rollback();
throw error;
}
}
case "datastore_listKinds": {
const { namespace } = args as any;
// Query the __kind__ special kind to get all entity kinds
const query = datastore.createQuery(namespace, '__kind__');
const [entities] = await datastore.runQuery(query);
const kinds = entities.map((entity: any) => {
const key = entity[datastore.KEY];
return key.name || key.id;
});
return {
content: [
{
type: "text",
text: JSON.stringify({
count: kinds.length,
kinds: kinds,
}, null, 2),
},
],
};
}
default:
throw new Error(`Unknown tool: ${name}`);
}
} catch (error: any) {
return {
content: [
{
type: "text",
text: JSON.stringify({
error: error.message,
details: error.stack,
}, null, 2),
},
],
isError: true,
};
}
});
// Start the server
async function main() {
const transport = new StdioServerTransport();
await server.connect(transport);
console.error("MCP Datastore Server running on stdio");
}
main().catch((error) => {
console.error("Fatal error in main():", error);
process.exit(1);
});
